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

gcl916

已于 2022-12-19 19:40:52 修改

阅读量2.4k

点赞数 1

文章标签：人工智能深度学习

于 2021-07-10 23:24:53 首次发布

本文链接：https://blog.youkuaiyun.com/gcl916/article/details/118642588

版权

这段代码展示了如何利用PyTorch和MobileNetV2模型，无需GPU，对个人图片进行分类。首先加载预训练的MobileNetV2模型，然后应用一系列图像预处理步骤，包括Resize、CenterCrop、ToTensor和Normalize。接着对图片进行前向传播计算，得到预测概率，并按降序排列，输出前五个最可能的类别及其置信度。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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