PyTorch 的“预期 CPU 张量（获得 CUDA 张量）”错误答案

【问题标题】："expected CPU tensor(got CUDA tensor)" error for PyTorchPyTorch 的“预期 CPU 张量（获得 CUDA 张量）”错误
【发布时间】：2017-09-01 09:00:04
【问题描述】：

在我使用经过训练的 PyTorch 模型的推理代码中，有什么问题？

出现运行时错误信息：“expected CPU tensor(got CUDA tensor)”

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()

我该如何解决这个问题？

【问题讨论】：

标签： python model inference pytorch tensor

【解决方案1】：

错误可能是cuda 中的模型不匹配以及您用作输入的变量 x，它是一个 CPU 张量。

尝试将.cuda() 添加到您的变量中，以便两者匹配：

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

【讨论】：

【解决方案2】：

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

此外，model.cuda() 和 x.cuda() 的行为略有不同：model.cuda() 会将 model 放在 GPU 上，但 x.cuda() 只是在 GPU 上返回一个新变量，而原始 x 保持不变。您必须明确地将返回值分配给x。

您可以找到详细讨论here。

【讨论】：