本系列文章配套代码获取有以下两种途径：
-
通过百度网盘获取：
链接：https://pan.baidu.com/s/1XuxKa9_G00NznvSK0cr5qw?pwd=mnsj提取码：mnsj
-
前往GitHub获取：
https://github.com/returu/PyTorch预训练模型：
-
torchvision.datasets：提供了一些加载数据的函数及常用的数据集接口，方便用户快速加载数据集； -
torchvision.models：包含了常用的模型结构（含预训练模型），如AlexNet、VGG、ResNet等，用户可以直接使用这些模型或者作为自己模型的起点进行微调； -
torchvision.transforms：提供了常用的图像预处理操作，如裁剪、旋转等。这些操作可以通过Compose类进行串联，方便用户构建自己的图像处理流程； -
torchvision.utils：包含其他一些有用的方法，如可视化工具等。
本次我们将使用torchvision.models来加载预训练模型，可以通过以下代码查看其包含的预定义模型：
import torchvision.models as models
dir(models)
# 使用类名创建模型实例
resnet = models.ResNet(blocks=[3, 4, 6, 3], num_classes=1000)# 使用便捷函数创建预训练模型实例
resnet18_pretrained = models.resnet18(pretrained=True)
使用resnet101完成图像识别分类：
下面，简要介绍下代码框架，具体内容可以查看文章配套的Notebook文件。
-
1、实例化一个resnet101模型：
# 使用resnet101来实例化一个具有101层的卷积神经网络
resnet = models.resnet101(pretrained=True)
# 查看网络结构信息，返回结果中的每一行就是深度学习中的层layers
resnet
ResNet(
(conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): Bottleneck(
(conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
......（略）......
)
)
(avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
(fc): Linear(in_features=2048, out_features=1000, bias=True)
)
-
2、加载预测图像并做预处理：
-
读取图像：
from PIL import Image
# 读取预测图像
img = Image.open('./image.jpg')
-
图像预处理：
# 定义图像预处理函数
from torchvision import transforms
preprocess = transforms.Compose([
transforms.Resize(256), # 将输入图像缩放到256*256大小
transforms.CenterCrop(224), # 围绕图像中心点将其裁剪为224*224大小
transforms.ToTensor(), # 将其转换为一个张量tensor
transforms.Normalize( # 对图像RGB分量进行归一化处理，使其具有定义的均值和方差
mean=[0.485,0.456,0.406],
std=[0.229,0.224,0.225]
)
])
img_t = preprocess(img)
-
按照网络输入要求对输入的张量进行维度扩展：
import torch
# 使用unsqueeze()函数进行维度扩展，维度变化为：[3, 224, 224] --> [1, 3, 224, 224]
batch_t = torch.unsqueeze(img_t , 0)
-
3、运行模型：
# 要进行推理，需要将网络置于eval模式
resnet.eval()
# 模型输出结果：一个由1000个分数的向量，每一个分数对应ImageNet中的每个类
out = resnet(batch_t)
out
输出结果：
tensor([[-3.2444e+00, -1.9064e+00, -3.7759e+00, -2.9615e+00, -3.5997e+00,
-2.9952e+00, -4.6614e+00, -2.3487e+00, -1.9136e+00, -2.0267e+00,
-7.4748e-01, -3.2986e+00, 3.9570e-01, -2.3369e+00, -1.6109e+00,
............
-2.1636e+00, 2.6281e-01, -2.6963e+00, 1.1257e+00, 3.9239e+00]],
grad_fn=<AddmmBackward0>)
-
4、获取模型预测结果标签及置信度：
-
读取文件生成标签列表：
该txt文件为ImageNet数据集包含的1000个标签。
with open('./imagenet_classes.txt') as f:
labels = [line.strip() for line in f.readlines()]
labels
输出结果：
['tench, Tinca tinca',
'goldfish, Carassius auratus',
............
'ear, spike, capitulum',
'toilet tissue, toilet paper, bathroom tissue']
-
获取最高分对应的索引：
通过max()函数获得输出结果中最高分数及其对应的索引。
_ , index = torch.max(out , 1)
-
使用softmax将输出结果归一化到[0,1]范围，可以认为是模型预测结果的置信度：
percentage = torch.nn.functional.softmax(out , dim=1)[0]*100-
获取预测结果标签及置信度：
labels[index[0]] , percentage[index[0]].item()
输出结果：
('tabby, tabby cat', 55.32239532470703)
-
获取预测结果前五名的标签及置信度：
_ , indices = torch.sort(out , descending=True)
n = 1
for idx in indices[0][0:5]:
print(f"预测第{n}名：类别为{labels[idx]},置信度为{percentage[idx].item()}")
n += 1
输出结果：
预测第1名：类别为tabby, tabby cat,置信度为55.32239532470703
预测第2名：类别为Egyptian cat,置信度为25.100072860717773
预测第3名：类别为tiger cat,置信度为17.506240844726562
预测第4名：类别为lynx, catamount,置信度为0.30555540323257446
预测第5名：类别为quilt, comforter, comfort, puff,置信度为0.09972941875457764
Reference： Antiga, Luca Pietro Giovanni, et al. Deep Learning with PyTorch. United States, Manning, 2020.
更多内容可以前往官网查看：
https://pytorch.org/
本篇文章来源于微信公众号: 码农设计师
