PyTorch搭建LeNet测试集实现

zhooooooou2024-03-08 10:25

测试集比较简单，直接上代码。

代码实现

python 复制代码

# 导包 不必多说
import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet

# 详细解释见下面
transform = transforms.Compose(
    [transforms.Resize((32, 32)),
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])


# 与训练集一样的分类
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# 模型实例化
net = LeNet()
net.load_state_dict(torch.load('LeNet.pth'))  # 载入在训练时保存的权重文件

im = Image.open('3.jpg')
im = transform(im)  # 预处理数据
im = torch.unsqueeze(im, dim=0)  # 将数据中增加一个batch维度

with torch.no_grad():
    outputs = net(im)
    # 寻找最大值所在的index索引值
    predict = torch.max(outputs, dim=1)[1].data.numpy()
# 最后打印预测结果
print(classes[int(predict)])

预处理数据函数

python 复制代码

transform = transforms.Compose(
    [transforms.Resize((32, 32)),
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

这一段是将数据预处理，相比于训练集中的数据预处理多了transforms.Resize((32, 32))，因为导入的图片尺寸不一定正确，需要先将图片的尺寸重新定义。

运行结果

我测试了飞机、汽车、鸟，飞机、汽车都可以识别出来。但鸟不行，可能是图片的像素太小，训练不到位。

把鸟给预测成猫了

总结

三天！从0开始，实现了LeNet。跟着b站上的视频，反复观看并记笔记，再自己手敲代码，再写出笔记。代码都能跑通实现，中间遇到的问题也靠自己独立解决了。对于自己来说还是比较有成就感的。但是我知道这点知识对于想要学好深度学习是远远远远不够的。还是要继续不断地学习。这样一篇一篇笔记也是我努力学习的见证！要努力成为很厉害的人！希望大家也是！