PyTorch的“期望的CPU张量（得到CUDA张量）”错误

在我训练PyTorch模型的推理代码中，出了什么问题？PyTorch的“期望的CPU张量（得到CUDA张量）”错误

有一个运行时错误信息： “预计CPU张量（GOT CUDA张量）”

import torch 
import torch.nn as nn 
#from __future__ import print_function 
import argparse 
from PIL import Image 
import torchvision.models as models 
import skimage.io 
from torch.autograd import Variable as V 
from torch.nn import functional as f 
from torchvision import transforms as trn 


# define image transformation 
centre_crop = trn.Compose([ 
trn.ToPILImage(), 
trn.Scale(256), 
trn.CenterCrop(224), 
trn.ToTensor(), 
trn.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 
]) 

filename=r'ed91.png' 
img = skimage.io.imread(filename) 

x = V(centre_crop(img).unsqueeze(0), volatile=True) 

model = models.__dict__['resnet18']() 
model = torch.nn.DataParallel(model).cuda() 

model = torch.load('mw_model0831.pth') 
#model.load_state_dict(checkpoint['state_dict']) 

#best_prec1 = checkpoint['best_prec1'] 
logit = model(x) 

print(logit) 
print(len(logit)) 
h_x = f.softmax(logit).data.squeeze()

我怎样才能解决这个问题？

来源

2017-09-01 제갈묘

可能错误是cuda 中的模型与您用作输入的变量x之间的不匹配，这是CPU张量。

尝试添加.cuda()到您的变量，都匹配：

x = V(centre_crop(img).unsqueeze(0), volatile=True).cuda()

来源

2017-09-01 13:50:44

的错误是因为model是GPU，而你的输入图像x是CPU。你必须确保它们都在GPU或CPU上。

此外，model.cuda()和x.cuda()表现略有不同：model.cuda()将放在GPU的model，但x.cuda()只返回一个新的变量上的GPU，离开原来的x不变。您必须明确指定返回值为x。

你可以找到详细的讨论here。

来源

2017-10-31 12:56:11 jdhao

PyTorch的“期望的CPU张量（得到CUDA张量）”错误

回答

相关问题