pytorch采坑(Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0)
问题描述
pytorch版本:2.0.1+cu117
使用pytorch的时候,我将模型部署到gpu上,然后将输入也部署到gpu上,但是却出现了如下报错
贴上我的源代码:
import torch.nn as nn
import torch.nn.functional as F
import torch
class Modeltest(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv1d(1, 20, 5)
self.conv2 = nn.Conv1d(20, 20, 5)
def forward(self, x):
x = F.relu(self.conv1(x))
return F.relu(self.conv2(x))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
input = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9,
10]).to(device).type(torch.FloatTensor)
input = torch.reshape(input, (1, -1))
modeltest = Modeltest().to(device)
output = modeltest(input)
output2 = output.sum() * 0.1
print(output.device)
print(output2.device)
从代码中可以看出,我定义了input,并写了to(device),然后对input进行了reshape操作,结果出现了报错
解决方案
我们最开始定义的input确实是在gpu上,但是使用了reshape之后,相当于重新定义了一个input,而默认情况下,新定义的变量都是在cpu上的,所以出现了问题,正确的代码应该如下:
import torch.nn as nn
import torch.nn.functional as F
import torch
class Modeltest(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv1d(1, 20, 5)
self.conv2 = nn.Conv1d(20, 20, 5)
def forward(self, x):
x = F.relu(self.conv1(x))
return F.relu(self.conv2(x))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
input = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9,
10]).type(torch.FloatTensor)
input = torch.reshape(input, (1, -1)).to(device)
modeltest = Modeltest().to(device)
output = modeltest(input)
output2 = output.sum() * 0.1
print(output.device)
print(output2.device)
或者第一个input定义的时候直接使用reshape,也可以,如下:
import torch.nn as nn
import torch.nn.functional as F
import torch
class Modeltest(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv1d(1, 20, 5)
self.conv2 = nn.Conv1d(20, 20, 5)
def forward(self, x):
x = F.relu(self.conv1(x))
return F.relu(self.conv2(x))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
input = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9,
10]).type(torch.FloatTensor).reshape(1, -1).to(device)
modeltest = Modeltest().to(device)
output = modeltest(input)
output2 = output.sum() * 0.1
print(output.device)
print(output2.device)