Lenet模型训练
- 数据集介绍
- 模型训练
- 数据集加载、预处理
- 定义模型网络结构
- 定义损失函数和优化器
- 模型训练
- 测试模型准确性
本文使用pytorch训练了一个对CIFAR-10数据集进行分类的Lenet模型,并使用Visdom可视化整个训练过程。
数据集介绍
CIFAR-10数据集共有60000张彩色图像,图像像素为32x32,分10类,每类6000张图片。其中50000张用于训练,另外10000张用于测试。测试数据集对每一类图像随机抽取1000张。
模型训练
本次训练过程遵循以下几个步骤:
- 数据集加载、预处理
- 定义模型网络结构
- 定义损失函数和优化器
- 训练模型并更新权重
- 测试模型准确率
数据集加载、预处理
使用torchvision下载并处理数据集
import torch
import torch.nn as nn
import torch.nn.functional as func
from torch import optim
import torchvision
import torchvision.transforms as transforms # 图像预处理包
from torchvision.transforms import ToPILImage # Tensor -> image,便于可视化
from visdom import Visdom # 训练过程可视化
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# 数据集预处理
# Compose组合多种预处理 [ToTensor, Normalize]
compose = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 下载训练数据集
train_set = torchvision.datasets.CIFAR10(root='./', train=True, download=True,transform=compose)
# 下载测试数据集
test_set = torchvision.datasets.CIFAR10(root='./', train=False, download=True, transform=compose)
# 加载数据集
# batch_size 为每次加载的样本个数
train_loader = torch.utils.data.DataLoader(train_set, batch_size=5, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=5, shuffle=True)
定义模型网络结构
定义网络时,需要继承nn.Moudle
,并实现它的forward
方法,把网络中具有可学习参数的层放在构造函数__init__
中,Lenet的网络结构这里不做过多分析,只进行网络实现:
# 定义网络结构
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
# 卷积层参数分别为:输入图片的通道数、输出通道数、卷积核size
self.conv1 = nn.Conv2d(3, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
# 全连接层: y = Wx + b
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = func.max_pool2d(func.relu(self.conv1(x)), 2)
x = func.max_pool2d(func.relu(self.conv2(x)), 2)
x = x.view(x.shape[0], -1)
x = func.relu(self.fc1(x))
x = func.relu(self.fc2(x))
x = self.fc3(x)
return x
net = Net()
# 查看模型结构
print(net)
# Net权重信息
params = net.parameters()
for name, parameters in net.named_parameters():
print(name, ":", parameters.size())
定义损失函数和优化器
nn实现了神经网络中的大多数损失函数,本文采用的损失函数为交叉熵损失函数nn.CrossEntropyLoss()
,并设置学习率为0.001,
反向传播计算完所有参数的梯度后,还需要使用优化方法更新网络的权重和参数,本文采用随机梯度下降法(SGD)的更新策略。
loss = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
模型训练
模型训练流程基本类似,不断地执行:
- 输入数据
- 前向传播+反向传播
- 更新参数
本文引入Visdom将训练过程可视化,Visdom的安装这里不赘述:
# visdom初始化
vis = Visdom()
vis.line([3.], [0.], win="Loss", opts=dict(title='train_loss'))
# 训练3个epoch,遍历完一遍数据集称为一个epoch
epoch = 3
# 模型训练
for each_epoch in range(epoch):
train_loss = 0.0
for i, data in enumerate(train_loader):
# 输入数据
inputs, labels = data
# 梯度清零
optimizer.zero_grad()
# forward + backward
out = net(inputs)
l = loss(out, labels)
l.backward()
# 更新参数
optimizer.step()
train_loss += l.data.item()
if i % 200 == 199:
print('[%d, %5d] loss: %.3f' % (each_epoch+1, i+1, train_loss/i))
vis.line([train_loss / i], [i + each_epoch * len(train_loader)], win='Loss', update='append')
# train_loss = 0
从图中大致能看到损失值的变化情况:
测试模型准确性
使用验证集在得到的模型上进行验证,准确率为57.48%
# 统计准确率
correct = 0
total = 0
for inputs, labels in test_loader:
outputs = net(inputs)
correct += (outputs.argmax(dim=1) == labels).float().sum().item()
total += labels.size(0)
print("accuracy: %f" % (correct / total))
结果:
... ...
[3, 9000] loss: 1.241
[3, 9200] loss: 1.240
[3, 9400] loss: 1.239
[3, 9600] loss: 1.238
[3, 9800] loss: 1.236
[3, 10000] loss: 1.235
Finished Training
accuracy: 0.574800