从经过训练的 Pytorch 模型获取预测

Question

I am using transfer learning to fine tune an inception_v3 model. After I train the model and store the best version off I am attempting to use it to generate predictions for my test set. Below is an example of my attempt on one image.

img_test=Image.open("img.png")

#Perform same transformations to image that the model used
transform_pipeline = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
img_test = transform_pipeline(img_test)

# I believe this is adding in the batch size of 1, but in looking around online it looked like I needed it
img = img_test.unsqueeze(0)
img = Variable(img)

    
model_ft(img)

When I do the above I get

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

Which seems to imply my model weights are on my gpu and the variable is on the cpu, how do I move one or the other over so I can use it, or reference one that is on the opposite processor?

Answer 1

As the error said, it seem that the input of the model (your img_test) is in the cpu.

Try to move the image to cuda before send it through your pre-trained model:

device = torch.device('cuda' if torch.cuda.is_available())
img_test = img_test.to(device)