[오류Error] RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=4, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(in_channels=4, out_channels=8, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()  # Flatten 모듈 추가
        self.fc1 = nn.Linear(8 * 16 * 16, 100)  # 입력 크기 조정 필요
        self.fc2 = nn.Linear(100, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = self.flatten(x)  # Flatten 함수 호출
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 모델의 정의 및 요약 정보 출력
model = SimpleCNN()
print(model)

from torchsummary import summary

INPUT_SIZE = 32
model = SimpleCNN() # 모델 인스턴스 생성
summary(model, input_size=(3, INPUT_SIZE, INPUT_SIZE))

 

간단한 딥러닝 코드를 돌리던 중 아래와 같은 에러가 발생했다. 

 

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-14-04359cf02c9c> in <cell line: 5>()
      3 INPUT_SIZE = 32
      4 model = SimpleCNN()  # 모델 인스턴스 생성
----> 5 summary(model, input_size=(3, INPUT_SIZE, INPUT_SIZE))

7 frames
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)
    454                             weight, bias, self.stride,
    455                             _pair(0), self.dilation, self.groups)
--> 456         return F.conv2d(input, weight, bias, self.stride,
    457                         self.padding, self.dilation, self.groups)
    458 

RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

 

위 오류 메시지는 모델의 입력 데이터와 가중치 데이터가 서로 다른 디바이스에 위치해 있기 때문에 발생합니다. 오류 메시지에는 `Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same`라고 나와 있어, 입력 데이터는 CUDA(즉, GPU)를 사용하는 반면, 모델의 가중치는 CPU에 있음을 의미합니다.

이 문제를 해결하기 위해 모델과 입력 데이터가 동일한 디바이스(모두 CPU 또는 모두 GPU)에 있어야 합니다. 만약 GPU에서 모델을 실행하고자 한다면, 모델과 입력 데이터 모두를 GPU로 옮겨야 합니다.

 

다음은 모델과 데이터를 GPU로 옮기는 코드입니다

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)  # 모델을 GPU로 옮김
summary(model, input_size=(3, INPUT_SIZE, INPUT_SIZE))

# 입력 데이터도 GPU로 옮겨야 할 경우, 다음과 같이 할 수 있습니다.
# data = data.to(device)

이 코드는 먼저 사용 가능한 경우 GPU를 사용하도록 설정하고, 모델을 해당 디바이스로 이동시킵니다. 

 

해결 완료!