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

本文主要是介绍pytorch预训练模型预测自己的图片，对大家解决编程问题具有一定的参考价值，需要的程序猿们随着小编来一起学习吧！

没有金贵的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]])

这篇关于pytorch预训练模型预测自己的图片的文章就介绍到这儿，希望我们推荐的文章对大家有所帮助，也希望大家多多支持为之网！