PyTorch에서 .to(device)를 사용하는 이유

안녕하세요, MoonLight입니다.

​

PyTorch를 사용하다보면, Tensorflow에는 없는 '.to(device)'라는 문구를 보게 됩니다.

​

이번 Post에서는 이 명령어가 무엇이고 왜 필요한지에 대해서 알아보도록 하겠습니다.

​

​

 

1. 딥러닝 연산과 하드웨어

딥러닝 모델은 엄청난 부동소수점 연산을 행렬 연산(Matrix Computation)을 수행합니다.

​

이러한 연산은 CPU에서도 가능하지만, GPU(Graphics Processing Unit)를 활용하면 훨씬 빠르게 처리할 수 있습니다.

​

​	CPU (중앙처리장치)	GPU (그래픽처리장치)
연산 방식	직렬(Serial) 연산	병렬(Parallel) 연산
핵심 개수	소수의 강력한 코어	수천수만 개의 작은 코어
연산 속도	일반적인 연산에 적합	대량의 연산을 빠르게 처리

 

 

 

2. PyTorch에서 .to(device)의 역할

PyTorch는 기본적으로 CPU에서 연산을 수행합니다.

​

하지만 GPU를 사용하려면, 데이터와 모델을 GPU 메모리로 옮겨야 합니다.

​

이를 위해 .to(device) 명령어를 사용합니다.

​

아래 Code는 .to(device)의 간단한 사용 예시입니다. 

import torch

# 사용할 장치(device) 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 텐서 생성
x = torch.randn(3, 3)

# CPU에서 연산
print(x.device)  # 출력: cpu

# GPU로 이동
x = x.to(device)
print(x.device)  # 출력: cuda (GPU가 있다면)

 

 

​

 

 

3. 모델과 데이터 모두 이동해야 하는 이유

PyTorch에서는 연산을 수행하는 모든 요소가 같은 장치에 있어야 합니다.

​

예를 들어, 모델이 GPU에 있고 데이터가 CPU에 있으면 연산할 수 없습니다.

​

따라서, 모델과 데이터를 동일한 장치로 옮겨야 합니다.

​

​

잘못된 코드 (CPU ↔ GPU 혼합)

x = torch.randn(3, 3).to("cuda")  # 데이터는 GPU에 있음
model = MyModel()  # 모델은 기본적으로 CPU에 있음

output = model(x)  # 오류 발생!

 

 

올바른 코드 (모델과 데이터 모두 GPU로 이동)

model = MyModel().to(device)  # 모델을 GPU로 이동
x = x.to(device)  # 데이터도 GPU로 이동

output = model(x)  # 정상 작동

 

 

 

​

​

4. GPU로 옮겨야 하는 요소들

모델을 GPU에서 학습하려면 기본적으로 학습 Data와 Model을 옮겨야 합니다.

이번에는 좀 더 구체적으로 이 내용에 대해서 알아보도록 하겠습니다.

​

​

 

Tensor (`.to(device)`)

모든 입력 데이터 및 라벨 Tensor는 GPU에서 연산을 수행해야 하므로 `.to(device)`를 사용하여 GPU로 이동해야 합니다.

import torch

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

x = torch.tensor([1.0, 2.0, 3.0])
x = x.to(device) # GPU로 이동

 

 

 

모델 (`model.to(device)`)

 

모델의 가중치(Weights)와 연산도 GPU에서 실행되도록 모델 자체를 GPU로 이동해야 합니다.

import torch.nn as nn

model = nn.Linear(10, 5) # 간단한 선형 모델
model.to(device) # 모델을 GPU로 이동

 

만약 모델을 GPU로 이동하지 않고 GPU Tensor와 연산하면 에러가 발생합니다.

 

 

 

 

옵티마이저(Optimizer)

 

PyTorch의 `torch.optim` 옵티마이저는 Tensor를 직접 다루지 않으므로 `.to(device)`를 사용하지 않아도 됩니다.

 

하지만 모델이 GPU로 이동하면 자동으로 옵티마이저도 GPU에서 작동합니다.

​

그러므로 별도로 `.to(device)`를 호출할 필요가 없습니다.

import torch.optim as optim

optimizer = optim.Adam(model.parameters(), lr=0.001) # 옵티마이저는 따로 이동 불필요

 

 

 

 

손실 함수 (Loss Function)

 

손실 함수 (`nn.CrossEntropyLoss()` 등) 자체는 GPU로 옮길 필요가 없지만, 계산할 때 입력 Tensor가 GPU에 있어야 합니다.

​

즉, GPU에서 연산이 수행될 수 있도록 손실 함수에 전달하는 데이터가 GPU에 있어야 합니다.

criterion = nn.CrossEntropyLoss() # 손실 함수 정의 (이동 불필요)

output = torch.tensor([[2.0, 1.0]], device=device) # GPU에 있는 데이터
target = torch.tensor([0], device=device) # GPU에 있는 라벨

loss = criterion(output, target) # 손실 함수 계산 (자동으로 GPU에서 실행)

 

 

 

 

 

데이터로더(DataLoader)에서 나온 배치 데이터

 

Dataloader를 통해 불러온 입력 Tensor와 라벨 Tensor도 GPU로 이동해야 합니다.

from torch.utils.data import DataLoader, TensorDataset

# 샘플 데이터 생성
x_data = torch.rand(100, 10)
y_data = torch.randint(0, 2, (100,))

dataset = TensorDataset(x_data, y_data)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

for x_batch, y_batch in dataloader:
    x_batch = x_batch.to(device)
    y_batch = y_batch.to(device)
    # 이제 GPU에서 학습 가능

 

 

 

정리

요소	GPU로 이동 필요 여부	이동 방법 (`.to(device)`)
Tensor (입력 데이터, 라벨)	필요	`tensor.to(device)`
모델(`nn.Module`)	필요	`model.to(device)`
옵티마이저 (`optim.Adam` 등)	불필요	X (자동 적용됨)
손실 함수 (`criterion`)	불필요	X (다만, 입력 Tensor는 GPU에 있어야 함)
DataLoader에서 나온 배치 데이터	필요	`batch.to(device)`

즉, 데이터(Tensor), 모델(nn.Module)은 반드시 `.to(device)`를 사용하여 GPU로 이동해야 합니다!

​

​

​

​

 

5. 주의할 점 : .to(device)는 새로운 객체를 반환한다는 사실

.to(device)는 기존 객체를 직접 변경하지 않고, 새로운 객체를 반환합니다. 따라서, 반환값을 저장해야 합니다.

x = torch.randn(3, 3)
x.to(device) # x는 여전히 CPU에 있음
print(x.device) # 출력: cpu

x = x.to(device) # 이제 x가 GPU에 있음
print(x.device) # 출력: cuda (GPU가 있다면)

 

 

 

 

6. Dataloader에 대한 고찰

Dataloader에서 나온 Batch Data는 Train Loop에서 일일이 옮겨야 하는 건가라는 의문이 들 수도 있습니다.

처음 Train Data를 읽을 때 한번에 모두 GPU로 옮기면 Train Loop에서 위와 같이 일일이 옮길 필요가 없을 것 같다는 생각이 들기도 합니다.

 

Dataloader에서 사용하는 Tensor Data가 이미 GPU에 옮겨져 있다면 안 옮겨도 되지 않을까요?

맞습니다. Dataloader에서 사용하는 `TensorDataset`의 Tensor들이 처음부터 GPU에 있으면, 이후 Train Loop에서 추가적으로 `.to(device)`를 호출할 필요가 없습니다.

​

하지만, 대부분의 경우 Train 데이터가 너무 크기 때문에 처음부터 GPU에 한 번에 올려놓는 것이 현실적으로 불가능합니다.

따라서 Batch 단위로 GPU로 옮기는 것이 가장 현실적인 방법이며, 이는 딥러닝에서 일반적으로 사용하는 방식입니다

​

​

​

​

도움이 되셨으면 좋겠습니다. 다음에 또 좋은 내용으로 찾아뵙겠습니다.

​

감사합니다.

​