num_workers & pin_memory in DataLoader

pytorch를 이용해 딥러닝 모델을 학습시킬 때 custom dataset을 이용할 경우

torch.utils.data.Dataset으로 데이터셋을 정의하고(input data type, augmentation 등)

torch.utils.data.DataLoader로 어떻게 데이터셋을 불러올 지(batch size, sampling 등) 정의한다.

 

학습을 시키다 보면 병목이 생기는 부분이 있는데 특히 데이터를 읽어서 가져올 때 시간이 오래 걸린다. 모델 학습을 하는데 시간을 써도 모자랄 판에 학습하기도 전에 불러오는 데이터에서 시간이 걸린다니... 즉, CPU를 이용해 데이터를 저장된 SSD나 HDD에서 읽어와 호스트의 메모리에 올리고 학습을 위해 GPU 메모리로 전달하는 과정에서 병목이 발생한다.

 

이를 좀 더 빠르게 수행할 수 있도록 파라미터를 설정할 수 있는데 대표적인 것이 바로 DataLoader에 있는 'num_workers'와 'pin_memory'다.

 

num_workers

pytorch 공식 문서에서 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)

즉, 서브 프로세스를 실행시켜 데이터를 불러오는데 그 때 실행시킬 서프 프로세스의 개수를 의미하는 파라미터다. 멀티프로세싱을 통해 저장된 데이터를 빠르게 가져오도록 하는 것이다.

 

그러면 많으면 많을수록 좋은거 아니야? 꼭 그렇지만은 않다. 자세한 내용은 아래 토론 글을 읽어보는 것을 추천한다.

 

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

Guidelines for assigning num_workers to DataLoader

num_workers가 몇 개가 적당한지에 대한 정답은 없다. 학습 환경의 GPU, CPU개수, I/O속도, 메모리 등과 같이 다양한 조건에 따라 최적의 값이 다르다고 한다. 아래 사진처럼 어떤 사람의 학습 환경에서는 num_workers가 5일 때 최적이지만 다른 사람의 환경에서는 다를 수 있다.

 

따라서 어떤 글들에서 보면 num_workers = GPU 개수 * 4 로 지정하면 좋다라는 말이 있는데 항상 그렇다고 생각하면 안된다는 것!

pin_memory

pin_memory (bool, optional) – If True, the data loader will copy Tensors into CUDA pinned memory before returning them. If your data elements are a custom type, or your collate_fn returns a batch that is a custom type, see the example below.

blog.naver.com/julie_eun1014/221116312880

CUDA 기반의 GPU Memory 이해하기

데이터를 CPU로 읽어들인 다음 GPU로 보내기 위해서는 GPU와 통신하기 위한 CPU의 메모리 공간이 필요하다. 이 때, 메모리를 할당시키는 기법을 memory pinning이라고 한다. memory pinning을 하는 데 사용되는 메모리의 종류는 pageable memory와 pinned memory가 있는데 pinned memory가 더 빠르다고 한다.

(pinned memory : Virtual memory를 사용하지 않는 Memory Mapped memory)

 

https://developer.nvidia.com/blog/how-optimize-data-transfers-cuda-cc/

즉, pinned memory는 위에 그림처럼 GPU에서 호스트에서 디바이스로 전송을 위한 staging area이고 pinned data transfer는 pinned memory와 pageable memory의 전송 비용을 줄이기 위해 데이터를 pin memory에 고정시켜 전송하는 방법이다.

 

우리가 pin_memory = True로 하게된다면 입력 데이터를 바로 pinned memory에 로드하여 빠르게 데이터 복사를 해 CUDA 연산을 효율적으로 할 수 있게 해준다. 따라서 시스템 메모리가 넉넉하다면 pin_memory = True로 하고 학습을 수행한다면 병목을 개선시킬 수 있다.