pytorch预训练模型预测自己的图片

没有金贵的GPU，用用巨头训练好的模型把自己的图片分分类。

其中的 dog.JPG是自己需要预测的图片。

imagenet_classes.txt在https://gist.github.com/ageitgey/4e1342c10a71981d0b491e1b8227328b

代码如下：

import torch
from PIL import Image
from torchvision import transforms
import torchvision.models as models

mobilenetV2 = models.mobilenet_v2(pretrained=True)
#print(mobilenetV2)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])

transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
            )])

img = Image.open("dog.JPG").convert('RGB')
print('Input image shape: ',img.size)

img_t = transform(img)
batch_t = torch.unsqueeze(img_t, 0)
mobilenetV2.eval()
out = mobilenetV2(batch_t)
print(out.shape)

with open('imagenet_classes.txt') as f:
    classes = [line.strip() for line in f.readlines()]

_, indices = torch.sort(out, descending=True)
percentage = torch.nn.functional.softmax(out, dim=1)[0] * 100

prediction = [[classes[idx], percentage[idx].item()] for idx in indices[0][:5]]

for i in prediction:
    print('Prediciton-> {:<25} Accuracy-> ({:.2f}%)'.format(i[0][5:], i[1]))