经典网络结构实现之LeNet-5
Pytorch实现手写数字识别
LeNet-5网络的简单实现
# 导入相应的包
import torch.nn as nn
# 搭建LeNet-5
class LeNet5(nn.Module):
def __init__(self):
super(LeNet5,self).__init__()
# 搭建卷积池化卷积池化层
self.conv1 = nn.Sequential(
nn.Conv2d( # (32*32*1)
in_channels = 1,
out_channels = 6,
kernel_size = 5,
stride = 1,
padding = 0
), # (28*28*6)
# 卷积和池化层之间添加ReLU激活函数(原文使用Sigmoid或者tanh函数)
nn.ReLU(),
nn.AvgPool2d(kernel_size = 2) # (14*14*6)
)
self.conv2 = nn.Sequential(
nn.Conv2d(6,16,5,1,0), # (10*10*16)
nn.ReLU(),
nn.AvgPool2d(2) # (5*5*16)
)
# 搭建全连接层
self.out = nn.Sequential(
nn.Linear(16*5*5,120),
nn.ReLU(), # 全连接层之间使用ReLU激活函数来进行激活
nn.Linear(120,84),
nn.ReLU(),
nn.Linear(84,10)
)
# 搭建前向传播接口
def forward(self,x):
x = self.conv1(x)
x = self.conv2(x)
# 将卷积层最后一个输出化成一个一维向量,然后才能传入全连接层
# 可以使用numpy中的reshape函数来进行替代
x = x.view(x.size(0),-1)
output = self.out(x)
return output
# look your network
myNet = LeNet5()
print(myNet)
LeNet-5应用于手写数字识别
Step1.导入所需的包
# 导包
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets,transforms
from torch.autograd import Variable
Step2.确定参数值
# 参数选取
lr = 0.01
momentum = 0.5
log_interval = 10 # 跑多少次batch进行一次日志记录
epochs = 10
batch_size = 64
test_batch_size = 1000
Step3.构建LeNet5网络模型
class LeNet(nn.Module):
def __init__(self):
super(LeNet,self).__init__()
self.conv1 = nn.Sequential( # (28*28*1)
nn.Conv2d(1,6,5,1,2), # (28*28*6)
nn.ReLU(),
nn.MaxPool2d(kernel_size = 2,stride = 2) # (14*14*6)
)
self.conv2 = nn.Sequential(
nn.Conv2d(6,16,5),
nn.ReLU(), # (10*10*16)
nn.MaxPool2d(2,2) # (5*5*16)
)
self.fc1 = nn.Sequential(
nn.Linear(16*5*5,120),
nn.ReLU()
)
self.fc2 = nn.Sequential(
nn.Linear(120,84),
nn.ReLU()
)
self.fc3 = nn.Linear(84,10)
def forward(self,x):
x = self.conv1(x)
x = self.conv2(x)
x = x.view(x.size()[0],-1)
x = self.fc1(x)
x = self.fc2(x)
x = self.fc3(x)
return x
Step4.定义每次训练的细节
def train(epoch): # 定义每个epoch的训练细节
model.train() # 设置为trainning模式
for batch_idx, (data, target) in enumerate(train_loader):
data = data.to(device)
target = target.to(device)
data, target = Variable(data), Variable(target) # 把数据转换成Variable
optimizer.zero_grad() # 优化器梯度初始化为零
output = model(data) # 把数据输入网络并得到输出,即进行前向传播
loss = F.cross_entropy(output,target) #交叉熵损失函数
loss.backward() # 反向传播梯度
optimizer.step() # 结束一次前传+反传之后,更新参数
if batch_idx % log_interval == 0: # 准备打印相关信息
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
Step5.定义测试函数
# 定义测试函数
def test():
model.eval() # 设置为test模式
test_loss = 0 # 初始化测试损失值为0
correct = 0 # 初始化预测正确的数据个数为0
for data, target in test_loader:
data = data.to(device)
target = target.to(device)
data, target = Variable(data), Variable(target) #计算前要把变量变成Variable形式,因为这样子才有梯度
output = model(data)
test_loss += F.cross_entropy(output, target, size_average=False).item() # sum up batch loss 把所有loss值进行累加
pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
correct += pred.eq(target.data.view_as(pred)).cpu().sum() # 对预测正确的数据个数进行累加
test_loss /= len(test_loader.dataset) # 因为把所有loss值进行过累加,所以最后要除以总得数据长度才得平均loss
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
Step6.训练预测模型
# 主函数启动训练
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #启用GPU
train_loader = torch.utils.data.DataLoader( # 加载训练数据
datasets.MNIST('../data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)) #数据集给出的均值和标准差系数,每个数据集都不同的,都数据集提供方给出的
])),
batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader( # 加载训练数据,详细用法参考我的Pytorch打怪路(一)系列-(1)
datasets.MNIST('../data', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)) #数据集给出的均值和标准差系数,每个数据集都不同的,都数据集提供方给出的
])),
batch_size=test_batch_size, shuffle=True)
model = LeNet() # 实例化一个网络对象
model = model.to(device)
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum) # 初始化优化器
for epoch in range(1, epochs + 1): # 以epoch为单位进行循环
train(epoch)
test()
torch.save(model, 'model.pth') #保存模型
Step7.手写数字识别测试
【注】使用画图板确保图片像素是28*28大小的,否则会报错
# 传入一张自己手写的数字进行识别
# 可将此主函数放入另一个python文件中
import torch
import cv2
import torch.nn.functional as F
# from modela import LeNet ##重要,虽然显示灰色(即在次代码中没用到),但若没有引入这个模型代码,加载模型时会找不到模型
from torch.autograd import Variable
from torchvision import datasets, transforms
import numpy as np
if __name__ =='__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.load('model.pth') #加载模型
model = model.to(device)
model.eval() #把模型转为test模式
img = cv2.imread("E:/6.jpg") #读取要预测的图片
trans = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)#图片转为灰度图,因为mnist数据集都是灰度图
img = trans(img)
img = img.to(device)
img = img.unsqueeze(0) #图片扩展多一维,因为输入到保存的模型中是4维的[batch_size,通道,长,宽],而普通图片只有三维,[通道,长,宽]
#扩展后,为[1,1,28,28]
output = model(img)
prob = F.softmax(output, dim=1)
prob = Variable(prob)
prob = prob.cpu().numpy() #用GPU的数据训练的模型保存的参数都是gpu形式的,要显示则先要转回cpu,再转回numpy模式
print(prob) #prob是10个分类的概率
pred = np.argmax(prob) #选出概率最大的一个
print(pred.item())
输入
预测输出
【注】该模型的预测准确度还是有待提升,待后续引入更好的训练模型进行训练