2017-09-01 4 views
1

Dans mon code d'inférence avec le modèle PyTorch formé, Quel est le problème?"tenseur de CPU attendu (obtenu tensor CUDA)" erreur pour PyTorch

Il y a un message d'erreur d'exécution: "tenseur CPU attendue (GOT tenseur 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() 

Comment puis-je résoudre ce problème?

Répondre

0

Probablement l'erreur est une discordance entre le modèle qui est en cuda et la variable x vous utilisez comme entrée, ce qui est un tenseur de CPU.

Essayez d'ajouter .cuda() à votre variable afin que les deux match:

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

L'erreur est parce que le model est sur GPU alors que votre image d'entrée x est sur CPU. Vous devez vous assurer qu'ils sont à la fois sur GPU ou CPU.

En outre, model.cuda() et x.cuda() se comporte un peu différemment: model.cuda() mettra le model sur GPU, mais x.cuda() juste retour d'une nouvelle variable sur GPU, laissant le x d'origine intacte. Vous devez affecter la valeur de retour à x explicitement. Vous pouvez trouver une discussion détaillée here.