用Pytorch搭建LeNet-5模型且基于CIFAR-10数据集训练
本文中的实例是Pytorch官方的一个demo，这里我结合最近正在看的书《Python深度学习 基于Pytorch》P125 Pytorch实现CIFAR-10多分类这节，进行了梳理与个人的总结，以加深自己的理解。
目的：
①了解LeNet-5模型
②加深Pytorch的学习
③从头至尾介绍搭建网络并进行训练预测的全过程
目录
-
- 简单介绍
-
- CIFAR-10数据集
- LeNet-5模型
- 查找Pytorch中API的具体内容
- 实例
-
- 1.构建神经网络
- 2.下载数据集并查看部分数据
- 3.训练模型
- 4.在测试集中随机预测四张图看看效果
- 5.测试模型在测试集上的准确率及10个类别的准确率
- 6.存储训练好的权重文件进行预测
- 整体源代码
简单介绍
CIFAR-10数据集
CIFAR-10数据集由10个类的60000个32x32彩色图像组成，以下是十个类别，每个类有6000个图像。有50000个训练图像和10000个测试图像。
classes = (‘plane’, ‘car’, ‘bird’, ‘cat’,‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’)
数据集分为五个训练批次和一个测试批次，每个批次有10000个图像。测试批次包含来自每个类的恰好1000个随机选择的图像。
训练批次包含随机顺序的图像，但一些训练批次可能包含来自一个种类的图像比另一个类更多。
总的训练批次包含来自每个类的正好5000张图像。
因为是彩色图像，3个通道，所以下面模型的输入为3×32×32
LeNet-5模型
1.模型架构
LeNet-5模型结构为输入层–卷积层–池化层–卷积层–池化层–全连接层–全连接层–输出，为串联模式，如上图所示
2.模型特点
a.每个卷积层包含3个部分：卷积、池化和非线性激活函数。
b.使用卷积提前空间特征。
c.采用降采样（Subsample）的平均池化层（Average Pooling）。
d.使用双曲正切（Tanh）的激活函数。
e.最后用MLP作为分类器。
3.模型的网络介绍
输入：
灰度图像，通道为1，尺寸为1×32×32
第一层：卷积层
LeNet-5模型接受的输入层大小是1×32x32。卷积层的过滤器的尺寸是5x5，深度（卷积核个数）为6，不使用全0填充，步长为1。则这一层的输出的尺寸为32-5+1=28，深度为6。本层的输出矩阵大小为6×28×28。
第二层：池化层
这一层的输入是第一层的输出，是一个6×28x28=4704的节点矩阵。本层采用的过滤器为2x2的大小，长和宽的步长均为2，所以本层的输出矩阵大小为6×14x14。
第三层：卷积层
本层的输入矩阵大小为6×14x14，使用的过滤器大小为5x5，深度为16。本层不使用全0填充，步长为1。本层的输出矩阵大小为16×10x10。
第四层：池化层
本层的输入矩阵大小是16×10x10，采用的过滤器大小是2x2，步长为2，本层的输出矩阵大小为16×5x5。
第五层：全连接层
本层的输入矩阵大小为16×5x5。将此矩阵中的节点拉成一个向量，那么这就和全连接层的输入一样了，本层的输出节点个数为120。
第六层：全连接层
本层的输入节点个数为120个，输出节点个数为84个。
第七层：全连接层
本层的输入节点为84个，输出节点个数为10个。
下面的神经网络模型是以LeNet-5为基准，其中就卷积核的个数不同，我看了书上和网上的模型，这里的参数都不同，猜测卷积核的个数可以自己定，但要保证各层之间维度的统一。
查找Pytorch中API的具体内容
方法一：在Pycharm中将鼠标移至函数名上，按住ctrl+鼠标左键
方法二：进入Pytorch官网——Docs——Search Docs中输入需要查找的API名称
实例
Pytorch中Tensor的通道排序：[batch，channel，height，width]
1.构建神经网络
定义一个LeNet()类，继承于nn.Module的父类，这个类中包括两个函数：初始化函数和前向传播函数。初始化函数中：模型使用到的网络层结构；前向传播函数：实现整个前向传播的过程。
class LeNet(nn.Module):
def __init__(self):
super(LeNet,self).__init__()
# 卷积层1：输入图像深度=3，输出图像深度=16，卷积核大小=5*5，卷积步长=1;16表示输出维度，也表示卷积核个数
self.conv1 = nn.Conv2d(in_channels=3,out_channels=16,kernel_size=5,stride=1)
# 池化层1：采用最大池化，区域集大小=2*2.池化步长=2
self.pool1 = nn.MaxPool2d(kernel_size=2,stride=2)
# 卷积层2
self.conv2 = nn.Conv2d(in_channels=16,out_channels=32,kernel_size=5,stride=1)
# 池化层2
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
# 全连接层1：输入大小=32*5*5，输出大小=120
self.fc1 = nn.Linear(32*5*5,120)
# 全连接层2
self.fc2 = nn.Linear(120,84)
# 全连接层3
self.fc3 = nn.Linear(84,10)
def forward(self,x):
x = F.relu(self.conv1(x)) # input(3, 32, 32) output(16, 28, 28)
x = self.pool1(x) # output(16, 14, 14)
x = F.relu(self.conv2(x)) # output(32, 10, 10)
x = self.pool2(x) # output(32, 5, 5)
x = x.view(-1, 32 * 5 * 5) # output(32*5*5)
x = F.relu(self.fc1(x)) # output(120)
x = F.relu(self.fc2(x)) # output(84)
x = self.fc3(x) # output(10)
return x
2.下载数据集并查看部分数据
import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim # 优化器
'''下载数据集'''
# transforms.Compose()函数将两个函数拼接起来。
# （ToTensor()：把一个PIL.Image转换成Tensor，Normalize()：标准化，即减均值，除以标准差）
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 训练集：下载CIFAR10数据集，如果没有事先下载该数据集，则将download参数改为True
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=False, transform=transform)
# 用DataLoader得到生成器，其中shuffle：是否将数据打乱；
# num_workers表示使用多进程加载的进程数，0代表不使用多进程
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,shuffle=True, num_workers=0)
# 测试集数据下载
testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,shuffle=False, num_workers=0)
classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# 显示图像
def imshow(img):
# 因为标准化normalize是：output = (input-0.5)/0.5
# 则反标准化unnormalize是：input = output*0.5 + 0.5
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
# transpose()会更改多维数组的轴的顺序
# Pytorch中是[channel，height，width],这里改为图像的[height，width，channel]
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()
# 随机获取部分训练数据
dataiter = iter(trainloader)
images, labels = dataiter.next()
# 显示图像
# torchvision.utils.make_grid()将多张图片拼接在一张图中
imshow(torchvision.utils.make_grid(images))
# 打印标签
# str.join(sequence)：用于将序列中的元素以指定的字符str连接成一个新的字符串。这里的str是' '，空格
# %5s：表示输出字符串至少5个字符，不够5个的话，左侧用空格补
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))
输出：
3.训练模型
'''训练模型'''
# 有GPU就用GPU跑，没有就用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = LeNet()
net=net.to(device)
# print("该网络共有 {} 个参数".format(sum(x.numel() for x in net.parameters())))
# # 该网络共有 121182 个参数
# 对于多分类问题，应该使用Softmax函数，这里CIFAR10数据集属于多分类，却使用了交叉熵损失函数，
# 是因为进入CrossEntropyLoss()函数内部就会发现其中包含了Softmax函数
loss_function = nn.CrossEntropyLoss() # 使用交叉熵损失函数
# 优化器选择Adam，学习率设为0.001
optimizer = optim.Adam(net.parameters(), lr=0.001)
# 打印查看神经网络的结构
# print(net)
# # LeNet(
# # (conv1): Conv2d(3, 16, kernel_size=(5, 5), stride=(1, 1))
# # (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
# # (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1))
# # (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
# # (fc1): Linear(in_features=800, out_features=120, bias=True)
# # (fc2): Linear(in_features=120, out_features=84, bias=True)
# # (fc3): Linear(in_features=84, out_features=10, bias=True)
# # )
for epoch in range(10): # 整个迭代10轮
running_loss = 0.0 # 初始化损失函数值loss=0
for i, data in enumerate(trainloader, start=0):
# 获取训练数据
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device) # 将数据及标签传入GPU/CPU
# 权重参数梯度清零
optimizer.zero_grad()
# 正向及反向传播
outputs = net(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
# 显示损失值
running_loss += loss.item()
if i % 2000 == 1999: # print every 2000 mini-batches
print('[%d, %5d] loss: %.3f' %(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
print('Finished Training')
输出：
[1, 2000] loss: 1.884
[1, 4000] loss: 1.602
[1, 6000] loss: 1.506
[1, 8000] loss: 1.445
[1, 10000] loss: 1.405
[1, 12000] loss: 1.377
[2, 2000] loss: 1.274
[2, 4000] loss: 1.248
[2, 6000] loss: 1.239
[2, 8000] loss: 1.226
[2, 10000] loss: 1.213
[2, 12000] loss: 1.185
[3, 2000] loss: 1.106
[3, 4000] loss: 1.124
[3, 6000] loss: 1.137
[3, 8000] loss: 1.109
[3, 10000] loss: 1.096
[3, 12000] loss: 1.089
[4, 2000] loss: 1.013
[4, 4000] loss: 1.015
[4, 6000] loss: 1.028
[4, 8000] loss: 1.037
[4, 10000] loss: 1.056
[4, 12000] loss: 1.009
[5, 2000] loss: 0.935
[5, 4000] loss: 0.964
[5, 6000] loss: 0.980
[5, 8000] loss: 0.962
[5, 10000] loss: 0.975
[5, 12000] loss: 0.993
[6, 2000] loss: 0.901
[6, 4000] loss: 0.908
[6, 6000] loss: 0.920
[6, 8000] loss: 0.924
[6, 10000] loss: 0.905
[6, 12000] loss: 0.944
[7, 2000] loss: 0.855
[7, 4000] loss: 0.864
[7, 6000] loss: 0.904
[7, 8000] loss: 0.896
[7, 10000] loss: 0.898
[7, 12000] loss: 0.890
[8, 2000] loss: 0.819
[8, 4000] loss: 0.828
[8, 6000] loss: 0.844
[8, 8000] loss: 0.880
[8, 10000] loss: 0.854
[8, 12000] loss: 0.863
[9, 2000] loss: 0.802
[9, 4000] loss: 0.795
[9, 6000] loss: 0.833
[9, 8000] loss: 0.848
[9, 10000] loss: 0.824
[9, 12000] loss: 0.834
[10, 2000] loss: 0.775
[10, 4000] loss: 0.758
[10, 6000] loss: 0.781
[10, 8000] loss: 0.806
[10, 10000] loss: 0.806
[10, 12000] loss: 0.829
Finished Training
为什么每计算一个batch，就需要调用一次optimizer.zero_grad() 将权重参数梯度清零？
是为了进行梯度累加。一定条件下，batchsize越大训练效果越好，梯度累加则实现了batchsize的变相扩大，如果accumulation_steps为8，则batchsize ‘变相’ 扩大了8倍。
通过这个特性能够变相实现一个很大batch数值的训练。
4.在测试集中随机预测四张图看看效果
'''预测四张图看看'''
dataiter = iter(testloader)
images, labels = dataiter.next()
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))
images, labels = images.to(device), labels.to(device)
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]for j in range(4)))
输出：
可以发现四张图只错了一张
5.测试模型在测试集上的准确率及10个类别的准确率
'''测试模型'''
correct = 0
total = 0
# with是一个上下文管理器
# with torch.no_grad()表示其包括的内容不需要计算梯度，也不会进行反向传播，节省内存
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(device), labels.to(device)
outputs = net(images)
# torch.max(outputs.data, 1)返回outputs每一行中最大值的那个元素，且返回其索引
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
# 打印10个分类的准确率
class_correct = list(0. for i in range(10)) #class_correct=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
class_total = list(0. for i in range(10)) #class_total=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(device), labels.to(device)
outputs = net(images) # outputs的维度是：4*10
# torch.max(outputs.data, 1)返回outputs每一行中最大值的那个元素，且返回其索引
# 此时predicted的维度是：4*1
_, predicted = torch.max(outputs, 1)
# 此时c的维度：4将预测值与实际标签进行比较，且进行降维
c = (predicted == labels).squeeze()
for i in range(4):
label = labels[i]
class_correct[label] += c[i].item()
class_total[label] += 1
for i in range(10):
print('Accuracy of %5s : %2d %%' % (classes[i], 100 * class_correct[i] / class_total[i]))
输出：
6.存储训练好的权重文件进行预测
torch.nn.Module模块中的state_dict变量存放训练过程中需要学习的权重和偏执系数，state_dict作为python的字典对象将每一层的参数映射成tensor张量
参考博客：state_dict详解
Python中.pth文件的作用
save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)
上面是已经训练好了的权重文件，下面随便用一张自己找的图片预测下
重新建一个py文件，读取训练好了的权重并对目标图像进行预测：
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('1.jpg')
im = transform(im) # [C, H, W]
# 输入pytorch网络中要求的格式是[batch，channel，height，width]，所以这里增加一个维度
im = torch.unsqueeze(im, dim=0) # [N, C, H, W]
with torch.no_grad():
outputs = net(im)
predict = torch.max(outputs, dim=1)[1].data.numpy() # 索引即classed中的类别
print(classes[int(predict)])
# 直接打印张量的预测结果
with torch.no_grad():
outputs = net(im)
predict = torch.softmax(outputs,dim=1) # [batch，channel，height，width],这里因为对batch不需要处理
print(predict)
输出：
整体源代码
main.py
import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim # 优化器
'''1.构建神经网络'''
class LeNet(nn.Module):
def __init__(self):
super(LeNet,self).__init__()
# 卷积层1：输入图像深度=3，输出图像深度=16，卷积核大小=5*5，卷积步长=1;16表示输出维度，也表示卷积核个数
self.conv1 = nn.Conv2d(in_channels=3,out_channels=16,kernel_size=5,stride=1)
# 池化层1：采用最大池化，区域集大小=2*2.池化步长=2
self.pool1 = nn.MaxPool2d(kernel_size=2,stride=2)
# 卷积层2
self.conv2 = nn.Conv2d(in_channels=16,out_channels=32,kernel_size=5,stride=1)
# 池化层2
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
# 全连接层1：输入大小=32*5*5，输出大小=120
self.fc1 = nn.Linear(32*5*5,120)
# 全连接层2
self.fc2 = nn.Linear(120,84)
# 全连接层3
self.fc3 = nn.Linear(84,10)
def forward(self,x):
x = F.relu(self.conv1(x)) # input(3, 32, 32) output(16, 28, 28)
x = self.pool1(x) # output(16, 14, 14)
x = F.relu(self.conv2(x)) # output(32, 10, 10)
x = self.pool2(x) # output(32, 5, 5)
x = x.view(-1, 32 * 5 * 5) # output(32*5*5)
x = F.relu(self.fc1(x)) # output(120)
x = F.relu(self.fc2(x)) # output(84)
x = self.fc3(x) # output(10)
return x
'''2.下载数据集并查看部分数据'''
# transforms.Compose()函数将两个函数拼接起来。
# （ToTensor()：把一个PIL.Image转换成Tensor，Normalize()：标准化，即减均值，除以标准差）
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 训练集：下载CIFAR10数据集，如果没有事先下载该数据集，则将download参数改为True
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=False, transform=transform)
# 用DataLoader得到生成器，其中shuffle：是否将数据打乱；
# num_workers表示使用多进程加载的进程数，0代表不使用多进程
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,shuffle=True, num_workers=0)
# 测试集数据下载
testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,shuffle=False, num_workers=0)
classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# 显示图像
def imshow(img):
# 因为标准化normalize是：output = (input-0.5)/0.5
# 则反标准化unnormalize是：input = output*0.5 + 0.5
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
# transpose()会更改多维数组的轴的顺序
# Pytorch中是[channel，height，width],这里改为图像的[height，width，channel]
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()
# 随机获取部分训练数据
dataiter = iter(trainloader)
images, labels = dataiter.next()
# 显示图像
# torchvision.utils.make_grid()将多张图片拼接在一张图中
imshow(torchvision.utils.make_grid(images))
# 打印标签
# str.join(sequence)：用于将序列中的元素以指定的字符str连接成一个新的字符串。这里的str是' '，空格
# %5s：表示输出字符串至少5个字符，不够5个的话，左侧用空格补
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))
'''3.训练模型'''
# 有GPU就用GPU跑，没有就用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = LeNet()
net=net.to(device)
# print("该网络共有 {} 个参数".format(sum(x.numel() for x in net.parameters())))
# # 该网络共有 121182 个参数
# 对于多分类问题，应该使用Softmax函数，这里CIFAR10数据集属于多分类，却使用了交叉熵损失函数，
# 是因为进入CrossEntropyLoss()函数内部就会发现其中包含了Softmax函数
loss_function = nn.CrossEntropyLoss() # 使用交叉熵损失函数
# 优化器选择Adam，学习率设为0.001
optimizer = optim.Adam(net.parameters(), lr=0.001)
# 打印查看神经网络的结构
# print(net)
# # LeNet(
# # (conv1): Conv2d(3, 16, kernel_size=(5, 5), stride=(1, 1))
# # (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
# # (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1))
# # (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
# # (fc1): Linear(in_features=800, out_features=120, bias=True)
# # (fc2): Linear(in_features=120, out_features=84, bias=True)
# # (fc3): Linear(in_features=84, out_features=10, bias=True)
# # )
for epoch in range(10): # 整个迭代10轮
running_loss = 0.0 # 初始化损失函数值loss=0
for i, data in enumerate(trainloader, start=0):
# 获取训练数据
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device) # 将数据及标签传入GPU/CPU
# 权重参数梯度清零
optimizer.zero_grad()
# 正向及反向传播
outputs = net(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
# 显示损失值
running_loss += loss.item()
if i % 2000 == 1999: # print every 2000 mini-batches
print('[%d, %5d] loss: %.3f' %(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
print('Finished Training')
'''4.在测试集中随机预测四张图看看效果'''
dataiter = iter(testloader)
images, labels = dataiter.next()
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))
images, labels = images.to(device), labels.to(device)
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]for j in range(4)))
'''5.测试模型在测试集上的准确率及10个类别的准确率'''
correct = 0
total = 0
# with是一个上下文管理器
# with torch.no_grad()表示其包括的内容不需要计算梯度，也不会进行反向传播，节省内存
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(device), labels.to(device)
outputs = net(images)
# torch.max(outputs.data, 1)返回outputs每一行中最大值的那个元素，且返回其索引
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
# 打印10个分类的准确率
class_correct = list(0. for i in range(10)) #class_correct=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
class_total = list(0. for i in range(10)) #class_total=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(device), labels.to(device)
outputs = net(images) # outputs的维度是：4*10
# torch.max(outputs.data, 1)返回outputs每一行中最大值的那个元素，且返回其索引
# 此时predicted的维度是：4*1
_, predicted = torch.max(outputs, 1)
# 此时c的维度：4将预测值与实际标签进行比较，且进行降维
c = (predicted == labels).squeeze()
for i in range(4):
label = labels[i]
class_correct[label] += c[i].item()
class_total[label] += 1
for i in range(10):
print('Accuracy of %5s : %2d %%' % (classes[i], 100 * class_correct[i] / class_total[i]))
'''6.存储训练好的权重文件'''
save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)
predict.py
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('1.jpg')
im = transform(im) # [C, H, W]
# 输入pytorch网络中要求的格式是[batch，channel，height，width]，所以这里增加一个维度
im = torch.unsqueeze(im, dim=0) # [N, C, H, W]
with torch.no_grad():
outputs = net(im)
predict = torch.max(outputs, dim=1)[1].data.numpy() # 索引即classed中的类别
print(classes[int(predict)])
# 直接打印张量的预测结果
with torch.no_grad():
outputs = net(im)
predict = torch.softmax(outputs,dim=1) # [batch，channel，height，width],这里因为对batch不需要处理
print(predict)