没有金贵的GPU,用用巨头训练好的模型把自己的图片分分类。
其中的 dog.JPG是自己需要预测的图片。
imagenet_classes.txt在https://gist.github.com/ageitgey/4e1342c10a71981d0b491e1b8227328b
代码如下:
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 mobilenetV2 = models.mobilenet_v2(pretrained=True) #print(mobilenetV2) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],#这是imagenet std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ #[1] transforms.Resize(256), #[2] transforms.CenterCrop(224), #[3] transforms.ToTensor(), #[4] transforms.Normalize( #[5] mean=[0.485, 0.456, 0.406], #[6] std=[0.229, 0.224, 0.225] #[7] )]) img = Image.open("dog.JPG") 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 print([(classes[idx], percentage[idx].item()) for idx in indices[0][:5]])