PyTroch DataLoader의 병렬 프로세싱!

다른 분이 작성한 PyTorch 코드를 보던 중, Dataset.__getitem__(index) 메서드에서 하드, 서버에 저장된 데이터를 불러오도록 구현한 것을 확인하였다.

__getitem__(index) 는 학습 과정에서 사용되는 매서드인데, 학습 과정에서 batch를 만들기위해 매번 하드에 접근한다면 큰 시간적 비효율이 있을텐데 왜 이렇게 다들 구현을 한건지 궁금해졌다.

 

그래서 찾아보던 중 PyTroch의 똑똑한 기능을 알게되어 공유해본다.

 
PyTrorch의 DataLoader는 기본적으로 병렬 프로세싱을 지원해 다음과 같이 CPU, GPU 자원을 시간효율적으로 사용할 수 있도록 지원한다. 여기서 병렬 프로세싱은, 아래처럼 GPU가 forward와 backward를 수행하는 동안, CPU는 디스크에서 데이터를 다음 배치에 쓰일 데이터를 미리 읽어오는 방식의 프로세싱을 말한다.

pytorch 병렬 프로세싱

DataLoader가 이처럼 똑똑하게 미리batch data를 구성하도록 하려면,

객체 정의 시 다음과같은 인자를 더 넣어주기만 하면 된다.

train_lodaer = DataLoader(train_dataset, batch_size=BATCH_SIZE,
                          num_workers=2,
                          pin_memory=True,
                          shuffle=True)

• num_workers: 병렬 데이터 로딩을 위해 사용될 추가 CPU 코어 개수. default는 0 이다.
• pin_memory: 사용 시 device(cuda)에 더 빠른 데이터 전송을 하게한다.
• shuffle: epoch단위로 데이터를 섞어줘, 모델이 같은 데이터 패턴을 만나지 않게 돼 잠재적 위험을 제거할 수 있다. 

물론 상황에 따라 유연하게 검증을 더 해봐야겠지만, 병렬 CPU코어를 사용하는 이상 우선은 __getitem__() 메서드를 작성할 때 너무 예민하게 생각하진 않아도 괜찮을 듯 하다..

(너무 많은 num_workers 배정 시 코드가 튕기니 조심하자.)


참고 URL https://wandb.ai/srishti-gureja-wandb/posts/reports/How-To-Eliminate-the-Data-Processing-Bottleneck-With-PyTorch--VmlldzoyNDMxNzM1