1.概述
LeNet在1998年发表,作为CNN的典型代表作,首先应用于手写数字识别任务中,并取得了很好的效果。
论文地址:http://yann.lecun.com/exdb/lenet/index.htmlhttp://yann.lecun.com/exdb/lenet/index.html
网络模型结构如下,由两个卷积层和三个全连接层组成。
2. pytorch实现LeNet5模型
2.1 模型文件 model.py
import torch.nn as nn
import torch.nn.functional as F
class LeNet_MNIST(nn.Module):
def __init__(self):
super(LeNet_MNIST, self).__init__()
self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5)
self.pool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5)
self.pool2 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(32 * 4 * 4, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool1(x)
x = F.relu(self.conv2(x))
x = self.pool2(x)
x = x.view(-1, 32 * 4 * 4)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
class LeNet_CIFAR(nn.Module):
def __init__(self):
super(LeNet_CIFAR, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5)
self.pool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5)
self.pool2 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(32 * 5 * 5, 120)
self.bn1 = nn.BatchNorm1d(120)
self.fc2 = nn.Linear(120, 84)
self.bn2 = nn.BatchNorm1d(84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool1(x)
x = F.relu(self.conv2(x))
x = self.pool2(x)
x = x.view(-1, 32 * 5 * 5)
x = F.relu(self.bn1(self.fc1(x)))
x = F.relu(self.bn2(self.fc2(x)))
x = self.fc3(x)
return x
其中分别实现了MNIST和CIFFAR10数据集的识别分类,由于MNIST是单通道的,CIFAR10是三通道的,故模型结构稍有不同。
2.2 模型训练代码 train.py
import sys
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from utils import train, test, show_result
from model import LeNet_CIFAR
sys.path.append('./../')
def main():
CIFAR_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))]
)
MNIST_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean=0.5, std=0.5)]
)
train_set = torchvision.datasets.CIFAR10(root='../data_set', train=True, download=True, transform=CIFAR_transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=128, shuffle=True, num_workers=0)
test_set = torchvision.datasets.CIFAR10(root='../data_set', train=False, download=True, transform=CIFAR_transform)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=128, shuffle=False, num_workers=0)
test_data_iter = iter(test_loader)
test_img, test_label = test_data_iter.next()
print(test_img.shape, test_label.shape)
net = LeNet_CIFAR().cuda()
loss_func = nn.CrossEntropyLoss().cuda()
optimizer = optim.Adam(net.parameters(), lr=0.0001)
loss_es = []
acc_es = []
best_acc = 0
for epoch in range(50):
loss = train(net, loss_func, optimizer, train_loader, epoch)
acc = test(net=net, test_loader=test_loader)
loss_es.append(loss)
acc_es.append(acc)
if acc > best_acc:
print('***************** save *********************')
best_acc = acc
torch.save(net, './LeNet-CIFAR.pt')
show_result(loss_es, acc_es, './LeNet5-CIFAI10.jpg')
if __name__ == '__main__':
main()
其中用到的工具文件 utils.py 如下,里面包含train和test函数以及绘制模型训练的损失曲线和准确率曲线,并保存训练结果图片到指定路径:
import torch
from matplotlib import pyplot as plt
import numpy as np
def train(net, loss_func, optimizer, train_loader, epoch):
"""
:param net:
:param loss_func:
:param optimizer:
:param train_loader:
:param epoch:
:return:
"""
running_loss = 0
total = 0
for step, (inputs, labels) in enumerate(train_loader):
inputs = inputs.cuda()
labels = labels.cuda()
out = net(inputs)
loss = loss_func(out, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
total += inputs.size(0)
running_loss /= total
print('Epoch: {}\tloss: {:.6f}\t'.format(epoch + 1, running_loss), end='')
return running_loss
def test(net, test_loader):
"""
:param net: if it is str, laod the pt, else is a model
:param test_loader:
:return:
"""
if isinstance(net, str):
model = torch.load(net)
else:
model = net
with torch.no_grad():
correct = 0
total = 0
for _, (inputs, labels) in enumerate(test_loader):
inputs = inputs.cuda()
labels = labels.cuda()
out = model(inputs)
predict_y = torch.max(out, dim=1)[1] # get max_value index
correct += torch.eq(predict_y, labels).sum().item()
total += inputs.size(0)
accuracy = correct / total
print('Accuracy: {:.3f}%.\t'.format(accuracy * 100))
return accuracy
def show_result(loss_es, acc_es, fig_save_name):
x = np.arange(len(loss_es))
plt.figure()
fig, ax1 = plt.subplots()
color = 'tab:red'
ax1.set_xlabel('epoch')
ax1.set_ylabel('loss', color=color)
ax1.plot(x, loss_es, label='loss', color=color)
ax1.tick_params(axis='y', labelcolor=color)
plt.legend()
ax2 = ax1.twinx() # 创建共用x轴的第二个y轴
color = 'tab:green'
ax2.set_ylabel('acc', color=color)
ax2.plot(x, acc_es, label='acc', color=color)
ax2.tick_params(axis='y', labelcolor=color)
fig.tight_layout()
plt.legend()
plt.savefig(fig_save_name, dpi=300)
3. 训练结果
经过50个epoch的训练后,对于MNIST数据集,可达到90多的准确率,对于CIFAR10可达到60多的准确率,值得注意的是,对于CIFAR10数据集,模型中加入BN层才会得到60多的准确率。训练的曲线如下图所示。