LeNet代码详解
model.py
import torch.nn as nn
import torch.nn.functional as F
class LeNet(nn.Module):#定义一个类,然后继承nn.Module这个父类
def __init__(self):#初始化函数
super(LeNet, self).__init__()#解决调用父类可能出现的问题
self.conv1 = nn.Conv2d(3, 16, 5)#nn.Conv2d来定义卷积层;(输入特征深度一共三个,卷积核个数,卷积核大小5*5)
self.pool1 = nn.MaxPool2d(2, 2)#定义了一个池化盒大小为2,步距也为2的最大池化操作(kernel_size,stride or kernel_size)只改变输出层的高和宽
self.conv2 = nn.Conv2d(16, 32, 5)#使用几个卷积核,就会生成深度为多少维的特征矩阵,通过第一次卷积以后在输入的特征层深度已经变成16了,然后采用32的卷积核数,大小依旧为5*5
self.pool2 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(32*5*5, 120)#全连接层,输入是一维的向量所以节点个数为32*5*5,输出的节点个数120是人为定义的
self.fc2 = nn.Linear(120, 84)#输入是上一层的输出
self.fc3 = nn.Linear(84, 10)#输入是上一层的输出,因为手写数字识别是一个十分类的输出所以输出节点个数定义为为10
def forward(self, x):#正向传播,x是输入x=[batch,channel,height,width]batch指一批图像的个数,如果我们将cifar10数据集按每批30个图片来分,那么batch=30,深度channel=3
x = F.relu(self.conv1(x)) # input(3, 32, 32)可以通过公式算出卷积以后的尺寸为28所以output(16,28,28)
x = self.pool1(x) # output(16, 14, 14)输出尺寸减半
x = F.relu(self.conv2(x)) # output(32, 10, 10)再次计算得(14-5+2*0)/1+1=10所以output(32,10,10)
x = self.pool2(x) # output(32, 5, 5)输出再次减半
x = x.view(-1, 32*5*5) # 使用view()方法进行展平因为batch没有,所以这里写-1,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
#先将conv1使用relu激活函数输出新的x,然后将x放入下采样层,然后再使用relu激活函数输出新的x,这时进入第二个下采样层,
# 接下来是要与全连接层进行拼接,所以先用view()方法进行展平,然后使用relu激活函数输出新的x进入第二个全连接层,
# 再relu激活函数输出新的x进入第三个全连接层,最后输出。
train.py
import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
def main():
transform = transforms.Compose(
#两个预处理方法
[transforms.ToTensor(),#可以将图片或一个numpy从H x W x C(每个维度像素值为[0, 255])变成C x H x W(每个维度像素值为[0.0, 1.0]),转变成一个Tensor。
#Normalize() 使用均值和标准差来归一化。
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 50000张训练图片
#先下载Cifar10数据集到train_set,root表示下载到当前目录下的data文件夹中去;
# train=true表示需要cifar10这个数据集;download第一次下载数据时改为true,下完以后改为false;
# transform是对图片进行预处理操作。
train_set = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
#将上面下载的图片集train_set分成每个batch为36的批次,shuffle=true表示分批时可以打乱。
train_loader = torch.utils.data.DataLoader(train_set, batch_size=36,
shuffle=True, num_workers=0)
# 10000张验证图片
# 第一次使用时要将download设置为True才会自动去下载数据集
val_set = torchvision.datasets.CIFAR10(root='./data', train=False,
download=False, transform=transform)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=5000,
shuffle=False, num_workers=0)
#使用iter()方法将它转换成迭代器,然后使用next()方法将迭代器中的图片、label取出来。
val_data_iter = iter(val_loader)
val_image, val_label = next(val_data_iter)
#因为是cifar10,也就是说这里有10类数据,我们定义一个classes元组,将10个种类都写出来:
# classes = ('plane', 'car', 'bird', 'cat',
# 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
#实例化LeNet网络
net = LeNet()
#并用CrossEntropyLoss()方法定义损失函数(这个方法中包括了softmax和loss函数)
loss_function = nn.CrossEntropyLoss()
#使用Adam优化器,并定义学习率lr = 0.001。
optimizer = optim.Adam(net.parameters(), lr=0.001)
#进入迭代训练过程
# range(5)意味着我们每轮要迭代5次
for epoch in range(5): # loop over the dataset multiple times
#定义runnning_loss用来累积损失
running_loss = 0.0
#然后再进入一个循环遍历训练集样本,不仅返回每一批训练完的数据data,还会返回这一批数据对应的步数step。
for step, data in enumerate(train_loader, start=0):
# get the inputs; data is a list of [inputs, labels]
#将data分为输入的图像inputs和对应的标签labels
inputs, labels = data
#利用zero_grad()函数将所有历史梯度清零,
# 每计算一个batch,就需要调用一次optimizer.zero_grad()函数清除历史梯度。
optimizer.zero_grad()
# forward + backward + optimize
#然后将输入的图片inputs放到网络中正向传播,返回预测标签放到outputs中
outputs = net(inputs)
#然后用预测值outputs和真实值labels放到损失函数中累计损失loss,然后用loss进行反向传播。
loss = loss_function(outputs, labels)
loss.backward()
#最后用optimizer.step()进行参数更新。
optimizer.step()
# print statistics
#然后每500步进行一次信息打印,在这个500步的训练中,我们将训练集传入net(),计算输出的预测标签predict_y与真实标签val_label相等的数量,计算准确率。
# 打印信息的时候就要打印迭代轮数、第多少步、平均损失、准确率。
running_loss += loss.item()#将每次产生的loss累计到running_loss中
if step % 500 == 499: #每隔500步打印一次训练的信息
with torch.no_grad():#with是一个上下文管理器 no_grad()函数可以保证不用每次都算每个节点的损失梯度
outputs = net(val_image) # [batch, 10] 将测试集的图片传入网络输出预测标签outputs
predict_y = torch.max(outputs, dim=1)[1]#找出 输出预测标签 与 第1维中输出的最大的index(第0维是batch,第1维是对应的10个节点也就是10种种类)
# 预测对了的标签的总数(因为是tensor格式所以要用item()方法转为数字)除以测试样本的数目就得到了准确率accuracy
accuracy = torch.eq(predict_y, val_label).sum().item() / val_label.size(0)
# 打印训练得到的参数,epoch表示训练到第几轮了;step表示在某一轮的第几步;然后是500步的平均误差;测试样本的准确率;
print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f' %
(epoch + 1, step + 1, running_loss / 500, accuracy))
running_loss = 0.0 #将误差清零,进行下一个500步的计算
print('Finished Training')
#训练结束,将网络中所有参数进行保存,保存路径为save_path。
save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)
if __name__ == '__main__':
main()
运行结果:
predict.py
输入的图像
#用来调用模型权重预测的文件
import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet
#1、在网上下载了一张飞机的图片,命名为1.jpg,放在项目文件夹下。
def main():
#2.导入相应的包之后,用Resize()方法将下载好的图片缩放处理成32X32的,然后转为向量并归一化。
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')
#3.实例化网络net,然后用load_state_dict()方法导入’Lenet.pth’文件;
net = LeNet()
net.load_state_dict(torch.load('Lenet.pth'))
#通过PIL模块载入1.jpg这张图像,使用transform()预处理,将[H,W,C]的图片处理成[C,H,W]的tensor,再用unsqueeze()方法加上N这个维度。
im = Image.open('1.jpg')
im = transform(im) # [C, H, W]
im = torch.unsqueeze(im, dim=0) # [N, C, H, W]
#将处理好的张量im传入网络,找到最大值输出中对应的index,根据index找到对应的种类名并输出。
#with torch.no_grad():
# outputs = net(im)#将处理好的张量im传入网络
# predict = torch.max(outputs, dim=1)[1].numpy()#找到最大值输出中对应的index
# print(classes[int(predict)])#根据index找到对应的种类名并输出。
with torch.no_grad():
outputs = net(im)#将处理好的张量im传入网络
predict = torch.softmax(outputs, dim=1)#找到最大值输出中对应的index
print(predict)#根据index找到对应的种类名并输出。
#用softmax函数可以将10个节点处输出的预测值打印出来,可以看到,我们输出第一个class的概率最高,也就是最可能是“plane”。
if __name__ == '__main__':
main()
运行结果:
softmax
max