pytorch中vgg的預訓練模型分類一張自己提供的圖片

import torch
import numpy
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
import torchvision.models as models

vgg = models.vgg16()
pre=torch.load('./vgg16-397923af.pth')
vgg.load_state_dict(pre)

r"""
vgg的pretrained模型是在imagenet上預訓練的，提供的是一個1000分類的輸出，每個類別標簽見：https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/c2c91c8e767d04621020c30ed31192724b863041/imagenet1000_clsid_to_human.txt
完美的圖片大小是224*224
transform就是三步走
unsqueeze后要加下划線才是原地操作，總是忘記
有了pth文件，要先torch.load一下，后load_state_dict
"""

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],#這是imagenet
                                 std=[0.229, 0.224, 0.225])

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


im='./1.jpeg'
im=Image.open(im)
im=tran(im)
im.unsqueeze_(dim=0)
print(im.shape)
# input()
out=vgg(im)
outnp=out.data[0]
ind=int(numpy.argmax(outnp))
print(ind)

from cls import d
print(d[ind])


print(out.shape)


# im.show()