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()
我怎样才能解决这个问题?
答
可能错误是cuda
中的模型与您用作输入的变量x
之间的不匹配,这是CPU张量。
尝试添加.cuda()
到您的变量,都匹配:
x = V(centre_crop(img).unsqueeze(0), volatile=True).cuda()
答
的错误是因为model
是GPU,而你的输入图像x
是CPU。你必须确保它们都在GPU或CPU上。
此外,model.cuda()
和x.cuda()
表现略有不同:model.cuda()
将放在GPU的model
,但x.cuda()
只返回一个新的变量上的GPU,离开原来的x
不变。您必须明确指定返回值为x
。
你可以找到详细的讨论here。