모델을 학습을 돌리는데, 훈련 중에 컴퓨터가 아예 멈추는 현상이 발생해 Dataloader pin memory과 관련이 있을 수 있어 정리해 보고자 한다.

개념

• pin_memory는 Pytorch의 DataLoader에서 사용되는 옵션이다.
• CPU 메모리에 있는 데이터를 GPU로 더 빠르게 전송할 수 있도록 도와주는 역할을 한다.
• pinned memory, 즉 고정된 메모리라는 뜻으로, CPU 메모리의 특정 공간을 잠궈서 이 데이터를 직접 GPU로 복사할 수 있게 해준다.

일반적으로 CPU 메모리는 비-고정 메모리라서 GPU로 복사할 때 추가적인 비용이 발생할 수 있다. pin_memory=True로 설정하면 메모리 복사 작업이 효율적으로 이루어 질 수 있게 된다.

간단한 비유를 하자면 이렇다.

데이터를 GPU로 보내는 과정을 물건을 배달하는 과정과 비교해보자.

• 일반 메모리(비-고정 메모리): 물건을 트럭에 실기 위해 창고에서 트럭앞으로 물건들을 가져오고, 다시 창고에 물건을 가지러 가는 동안, 트럭앞에 공간에서 다른 근무자가 다른 작업을 수행할 수 있다.

따라서 창고에 가질러 가는 시간이 소요되고, 다른 근무자가 작업을 하던 것을 정리하고 다시 물건을 트럭에 실어야 된다.

• 고정 메모리(pinned memory): 물건을 트럭에 실기 위해서 트럭 앞에 공간을 계속 점유하고 있어, 다른 근무자와의 접촉 없이 바로 트럭에 실으면 된다.

따라서 pin_memory=True로 설정하면 GPU로 데이터를 전송할 때 더 빨라진다고 이해할 수 있을 것이다.

from torch.utils.data import DataLoader, TensorDataset
import torch

data = torch.randn(100, 3)
labels = torch.randint(0, 2, (100,))

dataset = TensorDataset(data, labels)

dataloader = DataLoader(dataset, batch_size=10, shuffle=True, pin_memory=True)

왜 중요한가요?

장점

훈련 속도 : CPU에서 GPU로 데이터를 자주 전송해야 하기에 시간 단축이 된다. 또한 데이터가 클수록 전송하는 시간이 길어지는데, pin_memory는 이 시간을 줄여서 훈련을 더 빠르게 진행할 수 있도록 도와준다.

사용하면 좋은 경우

1. GPU를 사용해서 모델 훈련시킬 때.
2. 데이터가 많을 때.

단점

메모리 이슈: pin_memory를 사용하면 데이터를 고정된 메모리로 전송하여 유지하기 때문에 메모리를 더 많이 차지하게 된다. 고정 메모리는 메모리 공간을 확보한 후, 계속해서 다른 용도로 사용할 수 없기 때문에 자원이 제한적인 시스템에서는 메모리 부족 문제가 발생할 수 있다.

사용하면 안 좋은 경우

1. CPU만 사용할 때.
   CPU만 사용한다면 굳이 pin_memory를 사용할 필요가 없다. pin_memory는 GPU로 데이터를 전송할 때 속도를 높여주는 설정이니까 GPU를 사용하지 않는 경우에는 아무런 이점이 없는게 당연하다.

2. 메모리가 부족한 경우
   안그래도 부족한 메모리에 pin_memory까지 사용하면 내가 겪었던 것처럼 컴퓨터가 멈춰버릴 수도 있다.

3. 데이터셋이 작은 경우
   데이터셋이 작아서 CPU에서 GPU로 데이터를 전송하는 시간이 거의 발생하지 않는 경우에는 굳이 사용할 필요가 없다. 작은 데이터셋에서는 고정 메모리를 사용하는 것보다 기본 메모리 전송으로도 충분히 빠른 속도를 낼 수 있기 때문에 pin_memomry 설정을 했을 때 얻게 되는 성능 향상이 거의 없다.