본문 바로가기

AI/PyTorch

[PyTorch] DataLoader num_workers

 

PyTorch num_workers에 관하여

PyTorch에서의 데이터의 흐름과 적절한 num_workers값을 지정하는 방법

velog.io

  1. num_workers = 4 * num_GPU (or 8, 16, 2 * num_GPU)
  2. entry * batch_size * num_worker = num_GPU * GPU_throughtput
  3. num_workers = batch_size / num_GPU
  4. num_workers = batch_size / num_CPU

GPU에서 모델을 학습하는 속도와 CPU에서 DataLoader가 데이터를 준비해주는 속도는 다르다. 보통 CPU에서 수행되는 속도가 더 느리기 때문에 DataLoader는 멀티프로세싱을 통해 더 빠르게 GPU에 데이터를 조달한다. 이 때 사용할 멀티프로세스의 수가 num_workers이다.

num_workers는 적당한 값이 좋다. 너무 작으면 멀티프로세싱의 효과를 덜 보기 때문에 느릴 것이고, 너무 많으면 프로세스끼리 손발을 맞추기 위해 들어가는 오버헤드가 더 크기 때문에 오히려 느려진다.

 

 

DataLoader num_workers에 대한 고찰

Pytorch에서 학습 데이터를 읽어오는 용도로 사용되는 DataLoader는 torch 라이브러리를 import만 하면 쉽게 사용할 수 있어서 흔히 공식처럼 잘 쓰고 있습니다. 다음과 같이 같이 사용할 수 있겠네요. fr

jybaek.tistory.com

이렇듯 CPU에서의 작업을 빠르게 처리하고 task를 GPU로 던져서 GPU 사용률을 최대로 끌어내야 합니다. 그렇다면 CPU의 성능은 어떻게 이끌어내면 좋을까요? 기본적으로는 코드레벨에서 무언가 시도할 수 있겠지만 가장 단순한 방법은 작업을 단일코어가 아닌 멀티코어로 처리하는 것입니다. 1개 코어로 처리하고 있던 작업을 N개의 코어가 처리하게된다면? 당연한 이야기지만 훨씬 더 빠른 작업이 가능할겁니다. DataLoader에서 그것을 가능하게 해주는것이 바로 num_workers 파라미터 입니다.

 

그럼 처음 이야기한대로 데이터 프로세싱에 무조건 많은 CPU코어를 할당해주는 것이 좋은게 아닌가요? 꼭 그렇지는 않습니다. 코어 개수는 어차피 물리적으로 한정되어 있고 모든 코어를 전부 데이터 로드에 사용하게 된다면 다른 부가적인 처리에 딜레이가 생길수밖에 없습니다. 예를들면 데이터를 loading 하는 이외의 모든 작업이 영향을 받을 수 있겠죠. 그렇기 때문에 적당한 개수를 지정해줄 필요가 있습니다.

 

역시 적당히라는게 가장 어렵겠지만 하이퍼-파라미터를 튜닝하는 것처럼 결국 모델에 가장 적합한 num_workers 수치를 찾아내는 것도 파라미터 튜닝으로 볼 수 있습니다. num_workers 튜닝을 위해 고려해야 하는 것은 학습 환경의 GPU개수, CPU개수, I/O 속도, 메모리 등이 있습니다. I/O를 포함시킨 것은 데이터의 종류에 따라 디스크상에 존재하는 데이터를 로드하는것은 I/O에 상당히 많은 영향을 주고받을 수 있기 때문이고, 메모리는 loading된 데이터를 메모리상에 들고 있어야 하는 부담 때문에 포함되겠습니다.