Pytorch学习(四）—— Lenet模型训练

本文使用pytorch训练了一个对CIFAR-10数据集进行分类的Lenet模型，并使用Visdom可视化整个训练过程。

数据集介绍

CIFAR-10数据集共有60000张彩色图像，图像像素为32x32，分10类，每类6000张图片。其中50000张用于训练，另外10000张用于测试。测试数据集对每一类图像随机抽取1000张。

模型训练

本次训练过程遵循以下几个步骤：

1. 数据集加载、预处理
2. 定义模型网络结构
3. 定义损失函数和优化器
4. 训练模型并更新权重
5. 测试模型准确率

数据集加载、预处理

使用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