딥러닝 모델을 만들다 보면, GPU에서 학습을 하기 위해 모델을 GPU device에 할당하곤 합니다.
이때 두가지 방법으로 GPU에 모델을 올리는 방법과 차이점,
그리고 어떤 코드를 작성하면 좋을지에 대해 글을 작성 하였습니다.

개인적으로 학습하며 작성한 내용으로, 틀린 부분은 지적해 주시면 감사하겠습니다!

1. 모델을 GPU에 올리는 두 가지 방법

PyTorch를 사용하여 GPU에 모델을 할당할 때, 주로 다음의 코드를 이용합니다.

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class Net(nn.Module):
  def __init__(self):
    ...

  def forward(self, x):
    ...

net = Net()

# 1
net.cuda()

# 2
net = net.to(device)

그럼 위의 두 방법은 어떻게 다를까요?

기능적으로는 거의 동일하지만, 2번이 더 유연한 방법입니다.
왜냐하면 후자의 경우 cuda가 unavailable 하다면 'cpu'로 device가 초기화 되지만, 전자의 경우 AssertionError 를 발생시킵니다.

2. in-place (module, tensor)

위의 코드에서 후자의 경우를 추천했고, net = net.to(device) 와 같이 작성한 것을 볼 수 있습니다. 그럼 to(device)로 GPU에 할당한 모델을 리턴받아 다시 변수에 재할당 해줘야 하는 걸까요?

예시 코드를 작성하여 확인해보겠습니다.

class Net(nn.Module):
  def __init__(self):
    super().__init__()
    self.dummy_param = nn.Parameter(torch.empty(0))

  def forward(self, x):
    device = self.dummy_param.device

위와 같은 모듈에 대해 각 코드의 device 설정은 다음과 같습니다.

netA = Net()
print(netA.dummy_param.device)

netA.cuda()
print(netA.dummy_param.device)

> cpu
> cuda:0

netB = Net()
print(netB.dummy_param.device)

netB.to('cuda')
print(netB.dummy_param.device)

> cpu
> cuda:0

netC = Net()
print(netC.dummy_param.device)

netC = netC.to('cuda')
print(netC.dummy_param.device)

> cpu
> cuda:0

위와 같이 이상없이 GPU에 잘 할당된것을 볼 수 있습니다.

하지만, Tensor에서 재밌는 모습을 볼 수 있습니다.

x = torch.randn(1, device='cpu')
print(x.device)

x.cuda()
print(x.device)

x = x.cuda()
print(x.device)

> cpu
> cpu
> cuda:0

x = torch.randn(1, device='cpu')
print(x.device)

x.to('cuda')
print(x.device)

x = x.to('cuda')
print(x.device)

> cpu
> cpu
> cuda:0

이렇게 Tensor를 GPU에 할당하는 경우 inplace 옵션이 적용되지 않음을 볼 수 있습니다.

결론

device만 신경 쓴다면 to(device)나 cuda()나 동일하지만,
to(device) 방식은 dtype이나, memory-layout도 변경해준다고 합니다.
따라서 좀 더 기능적으로 사용하기 위해 to(device) 방식을 사용하는 것이 더 좋아 보입니다.
또한 예상치 못한 에러를 방지하기 위해 GPU에 올린 객체들을 변수에 재할당하여 사용할 수 있도록 습관을 들이는 것이 좋습니다.

reference.

https://discuss.pytorch.org/t/what-is-the-difference-between-doing-net-cuda-vs-net-to-device/69278/9

https://stackoverflow.com/questions/58926054/how-to-get-the-device-type-of-a-pytorch-module-conveniently