使用torch训练模型时,出现错误: “Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same”
错误原因: 模型放入GPU后也只能处理数据类型为cuda的GPU数据,但你在把数据输入进模型之前并未做转换,依然是张量类型 。因此出错,二者必须一致。
例如:
device = torch.device('cuda:0')
Mymodel = VGGModel().to(device)
...
...
for step, (data, label) in enumerate(train_loader):
model.train()
output = model(data).to(device)
问题就在output = model(data).to(device),虽然我们把模型用.to(devce)传入了GPU,但data和label两个张量并没有变为GPU能处理的cuda类型数据,仍然是张量数据,此时model因为用了.to(device) 后只能处理cuda类型数据,当然二者不能兼容,因此报错。
解决方法: 在读出data和label后使用.to(device)转换成cuda数据后在输入进模型。 代码改进如下:
device = torch.device('cuda:0')
Mymodel = VGGModel().to(device)
...
...
for step, (data, label) in enumerate(train_loader):
data, label= data.to(device), label.to(device)
model.train()
output = model(data).to(device)
问题解决。
|