DataLoader num_workers에 대한 고찰

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

from torch.utils.data import DataLoader

상세한 설명이 기술되어 있는 공식 문서는 아래 링크에서 살펴볼 수 있습니다.

https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader

인자로 여러가지 파라미터를 넘길수 있는데 여기서 이야기하고자 하는 부분은 num_workers인데 공식문서의 설명은 다음과 같이 되어 있습니다.

num_workers (int, optional) – how many subprocesses to use for data loading. 
      0 means that the data will be loaded in the main process. (default: 0)

사실 읽어봐도 감이 잘 안옵니다. default 값은 0인데 data 로딩을 위해 몇 개의 서브 프로세스를 사용할 것인지를 결정한다는 이야긴데 결국 데이터 로드 멀티 프로세싱에 대한 이야기입니다. data loading이라면 그냥 잔뜩 많이 사용하는게 좋은게 아닌가? 라고 생각할 수 있지만 여기에는 살짝 미묘한 부분이 있습니다. CPU와 GPU 작업간의 밸런스인데요. 이제부터 하나씩 이야기해보도록 합시다.

머신러닝에서는 (엄청나게 많은) 단순한 행렬연산을 GPU를 통해 빠르게 처리하는데 우리가 비싼 그래픽카드를 사놓고 제대로 일을 시키고 있지 않다면 그것만큼 슬픈일은 없을겁니다. 아래 그림을 살펴볼텐데 CPU에서 작업을 GPU로 넘기기 위해 데이터를 전처리하는 과정(아래 그림 빨간색 선)이 너무 오래 걸린다면 GPU가 그만큼 일을 하지 않게된다는 것을 의미합니다.

한편 빠른 전처리(위 그림 보라색 선)를 통해 CPU가 task를 바로바로 GPU로 던져줄 수 있다면 GPU는 쉬는시간 없이 계속 일을 하게 될겁니다. 머신러닝에서 가장 많은 시간을 소비하게 되는 구간이 GPU라는 것을 생각해봤을때 GPU는 놀면 안되겠죠. (일반적인 머신러닝 상황입니다. GPU, 모델의 종류 등에 따라 예외적인 상황이 있습니다)

GPU를 잘 활용하는 좋은 예를 가져왔습니다. 아래 첨부된 이미지에서 GPU 사용량(GPU-Util)을 살펴보세요.

nvidia-smi

이렇듯 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된 데이터를 메모리상에 들고 있어야 하는 부담 때문에 포함되겠습니다.

이런 여러가지 이슈들 때문에 num_workers 값 튜닝에 대해서 토론까지 진행을 하기도 합니다. 관련된 토론내용은 아래 링크에서 확인하실 수 있습니다.

https://discuss.pytorch.org/t/guidelines-for-assigning-num-workers-to-dataloader/813

Guidelines for assigning num_workers to DataLoader

위에 토론에는 생각해볼만한 다양한 이슈들을 확인할 수 있기 때문에 일독을 권합니다.

여기까지 num_workers 파라미터가 어떤 역할을 수행하며 어떻게 값을 세팅하면 좋을지에 대해서 이야기를 해봤는데 결국 최종 선택값은 사용자의 몫이겠습니다. 보통의 일반적인 환경에서 오픈소스로 풀려있는 모델을 학습시킬때는 코어 개수의 절반정도 수치면 무난하게 시스템 리소스를 사용하며 학습이 가능했습니다. 다시 말하지만 최종 선택은 사용자 본인 입니다. 참고만 하시길.  :-)