pytorch深度学习实战——预训练网络

CheneyTAT

已于 2022-03-26 21:26:48 修改

阅读量2.9k

点赞数

分类专栏：笨方法学python3 文章标签： python 开发语言

于 2022-03-20 19:08:45 首次发布

本文链接：https://blog.csdn.net/weixin_66648679/article/details/123618857

版权

来源：《Pytorch深度学习实战》，2.1，一个识别图像主体的预训练网络

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

dir(models)
alexnet = models.AlexNet()
resnet = models.resnet101(pretrained=True)
resnet

preprocess = 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("D:\\PyCharm\\pythonProject\\bobby.PNG")
img.show()
img_t = preprocess(img)
batch_t = torch.unsqueeze(img_t, 0)
resnet.eval()
out = resnet(batch_t)
out

with open("D:\\PyCharm\\pythonProject\\imagenet_classes.txt") as f:
    labels = [line.strip() for line in f.readline()]

_, index = torch.max(out, 1)

percentage = torch.nn.functional.softmax(out, dim=1)[0] * 100
labels[index[0]], percentage[index[0]].item()