라임오렌지파이와 일상

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를 설정하려고 하니, 이런 문제가 발생합니다.

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

• Dataloader에서의 num_workers 설정의 필요성: https://jybaek.tistory.com/799
• PyTorch Windows FAQ: https://pytorch.org/docs/stable/notes/windows.html#multiprocessing-error-without-if-clause-protection

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

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

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