CNN卷积神经网络03--LeNet推理

目录

推理阶段

获取图片

图片预处理

从网络前向传播的数组获取结果

完整代码

---

推理阶段

训练完模型保存后,在推理阶段只需要实例化模型,加载参数,然后对图片进行预处理即可获取推理结果,实际上这部分和训练时候进行评估的代码十分类似

获取图片

直接从CIFAR10 的测试集中选6张图片

图片预处理

在训练阶段,我们对图片进行了归一化及标准化操作以及调整维度顺序hwc-->chw,预测的时候,我们需要对待预测的图片进行同样的处理,并且进行升维组成BCHW格式的数组

img_data = np.transpose(((np.array(img) / 255) - mean) / std, (2, 0, 1))[None, ...]
img_data = torch.tensor(img_data, dtype=torch.float32)

从网络前向传播的数组获取结果

网络的前向传播的到的数组实际还不是我们真正理解的结果

根据网络的定义,前向传播返回的结果是[B,10]的数组,

需要把这个数组解释成人类理解的结果,对数组进行sorftmax函数处理以获取我们对其输出的各类的置信度信息,而最大值的索引即为预测的label

img_data = np.transpose(((np.array(img) / 255) - mean) / std, (2, 0, 1))[None, ...]
img_data = torch.tensor(img_data, dtype=torch.float32)
pred = model(img_data)
conf = torch.softmax(pred, dim=1)
conf, index = conf.max(dim=1)

完整代码

import numpy as np
import torch
import torchvision
from matplotlib import pyplot as plt

from net import LeNet


def get_test_img(n):
    test_set = torchvision.datasets.CIFAR10(
        root=r'D:\Cproject\ex001\CIFAR10',  # 如果下载,下载的位置,也可以手动下载,放到这个位置来
        train=False,
        download=False,  # 下载地址  https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
        transform=None
    )
    samples = []
    for i in range(n):
        img, label = test_set[i]
        samples.append(([img, label]))
    return samples


def main():
    fig = plt.figure()  # 初始化一张画布
    samples = get_test_img(6)
    classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
    target = []
    for i in range(len(samples)):
        plt.subplot(2, 3, i + 1)
        plt.imshow(samples[i][0])
        plt.title(classes[samples[i][1]])
        target.append(classes[samples[i][1]])
    # 展示一下原始图片
    plt.show()
    model = LeNet(10)
    model.load_state_dict(torch.load('model.pth', map_location='cpu'))
    model.eval()
    with torch.no_grad():
        mean = (0.5, 0.5, 0.5)
        std = (0.5, 0.5, 0.5)
        result = []
        for img, label in samples:
            img_data = np.transpose(((np.array(img) / 255) - mean) / std, (2, 0, 1))[None, ...]
            img_data = torch.tensor(img_data, dtype=torch.float32)
            pred = model(img_data)
            conf = torch.softmax(pred, dim=1)
            conf, index = conf.max(dim=1)
            result.append([conf.item(), index])

    for i in range(len(samples)):
        print(f"""{i:<3}  label:{target[i]:<10}  pred:{classes[result[i][1]]:<10}  conf:{result[i][0]:.3f}""")
    return


if __name__ == '__main__':
    main()

 结果如下

0    label:cat         pred:cat         conf:0.684
1    label:ship        pred:ship        conf:0.994
2    label:ship        pred:airplane    conf:0.214
3    label:airplane    pred:ship        conf:0.634
4    label:frog        pred:frog        conf:0.473
5    label:frog        pred:frog        conf:0.714