2017-09-01 13 views
1

訓練されたPyTorchモデルの私の推論コードで何が問題なのですか?PyTorchの "予想されるCPUテンソル(CUDAテンソルを得ました)"エラー

ランタイムエラーメッセージがあります: "予想されるCPUテンソル(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() 

私はこれをどのように修正することができますか?

答えて

0

cuda のモデルと、入力として使用している変数x(CPUテンソル)の間のミスマッチである可能性があります。

は、あなたの変数に.cuda()を追加しようとするので、両者が一致:あなたの入力画像xがCPU上にある間modelはGPU上にあるため

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

エラーです。それらがGPUかCPUの両方であることを確認する必要があります。

はまた、model.cuda()x.cuda()は少し異なる動作をします:model.cuda()は、GPU上model入れますが、x.cuda()だけのオリジナルxはそのまま残し、GPU上で新しい変数を返します。戻り値を明示的にxに割り当てる必要があります。

詳しくは、hereをご覧ください。

関連する問題