[PyTorch] GPU을 사용할 때 to(device)와 cuda() 차이

PyTorch에서는 GPU을 사용하여 모델의 트레이닝 속도를 높일 수 있다.

모델을 트레이닝 하는데 GPU을 사용하려면 데이터와 모델을 GPU에 할당해야한다.

이때 텐서나 모델을 GPU에 올리는 두 가지 방법이 있는데 이번 글에서는 두 가지 방법과 그 차이점을 살펴보자.

---

GPU을 사용하는 두 가지 방법

PyTorch에서 텐서나 모델을 GPU에 할당할 때 두 가지 방법이 있는데 .to(device)와 .cuda()이다.

먼적 GPU에 올리는 방법을 각각 알아보자.

 

먼저 .to(device)은 CPU, GPU 및 TPU를 포함한 모든 장치로 텐서나 모델을 이동할 수 있는 일반적인 방법이다.

보통 아래와 같은 방법으로 사용한다.

import torch

# Create a tensor on the CPU
x = torch.randn(10, 10)

# Move the tensor to a GPU device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x = x.to(device) # or x.to('cuda')

# Move the tensor back to CPU
x = x.to('cpu')

 

한편 .cuda()는 GPU로 옮기는 보다 구체적인 방법으로 이 방법을 사용하면 텐서를 GPU로 이동하거나 사용 가능한 GPU가 없으면 AssertionError에러를 뱉는다.

이 방법은 PyTorch의 빌드가 CUDA을 지원하는 경우에만 사용할 수 있다.

예시를 통해 방법을 살펴보자.

import torch

# Create a tensor on the CPU
x = torch.randn(10, 10)

# Move the tensor to the default GPU
x = x.cuda()

# Move the tensor back to CPU
x = x.cpu()

두 방법 모두 동일한 결과를 가져오기는 하지만 .to(device)을 사용할 때 코드를 변경하지 않고 텐서를 다른 장치로 이동시킬 수 있으므로 좀더 유연하다고 할 수 있다.

---

to(device)와 cuda()의 차이

가장 큰 차이점은 .to(device)은 모델을 CPU나 GPU 둘 다 이동할 수 있는 반면 .cuda()은 CUDA을 지원하는 GPU로만 모델을 이동할 수 있다.

우리가 가장 관심있을 속도 측면에서 두 방법이 큰 차이는 없지만 .cuda()가 아주 조금 더 빠를 수는 있기는 하지만 CPU만 사용하는 환경에서는 코드를 뜯어고쳐야 하는 불편함이 있다.

하지만 .to(device)의 경우 위에서 살펴본 것처럼 코드를 변경하지 않고도 다양한 디바이스에서 코드를 실행할 수 있는 장점이 있다.

 

결국 사용자의 하드웨어 설정에 따라 선택을 해야한다.

만약 GPU가 있고 병렬 처리 기능을 잘 활용하고 싶다면 .cuda()가 좋은 옵션이 될 것이다.

하지만 디바이스를 유연하게 설정하고 코드가 다른 유형도 지원하게 하고 싶다면.to(device)을 선택하는 것이 좋을 것이다.