[PyTorch] Windows에서 num_workers로 데이터로더 성능 높이기

RuntimeError: cuda runtime error (801): operation not supported at ..\torch/csrc/generic/StorageSharing.cpp:247

num_workers와 DataLoader에서 GPU 복사를 구현할 경우 발생하는 문제

기존 PyTorch로 학습을 진행하다가, single-thread로 동작하는 DataLoader의 bottleneck때문에 GPU를 최대한 활용하지 못하는 문제점이 있었습니다. 다만, PyTorch에서 기본적으로 제공하는 DataLoader의 num_workers를 설정하려고 하니, 이런 문제가 발생합니다.

데이터로더에서 이미 CUDA로 복사한 데이터를 넘겨주게 되면, Multiprocessing에 문제가 발생합니다.

각종 레퍼런스들을 찾다가, 이러한 내용을 발견했습니다.

그래서 생각을 해보니, Dataset에서는 CPU Tensor를 전달해주고, DataLoader를 순회하는 Training Time에 데이터를 CUDA에 복사하면 될 것 같았습니다. 그래서, 해당 방식을 이용해서 Dataset에는 CPU만을 이용하여 데이터를 로드하도록 해보았습니다.

상단의 데이터셋은 제가 Custom으로 작성한 torch.utils.data.Dataset의 상속 클래스입니다. OpenCV를 이용해 데이터를 CPU로만 로드하도록 하였습니다. 그리고 난 뒤, 아래와 같이 데이터로더 순회 후에 모델의 입력 직전에 GPU로 복사하도록 했습니다.

위 방법을 쓰니, 더이상 num_workers에 관련한 이슈는 발생하지 않아서 좋은 해결법이 되었습니다….!!

답글 남기기

이메일 주소는 공개되지 않습니다. 필수 필드는 *로 표시됩니다