【Pytorch基础学习】（4）模型训练与验证

本节总结本系列的前面几节，以LeNet神经网络为例，实现一套模型训练与验证的完整流程，代码如下，每一步详情请见代码注释：

import torch
import torch.nn as nn
import torchvision
import numpy
import matplotlib.pyplot as plt
import datetime
import os
from tqdm import tqdm

# 设置设备的选取，如果本机配置了GPU、CUDA、CUDNN即选用CUDA，否则选用CPU
device = (torch.device('cuda') if torch.cuda.is_available()
          else torch.divide('cpu'))
print(f"Training on device {device}")

# 设置数据的地址以及模型权重的保存地址
data_path = r'E:\code\ML\ML\practice\data\Fashion_MNIST'
pt_path = r'E:\CS\code\ML\ML\practice\model\lenet\pt\\'

# 加载训练集（关于数据集的加载以及本地数据集的创建可以参考本系列之前的内容）
MNIST = torchvision.datasets.MNIST(data_path, train=True, download=True,
                                          transform=torchvision.transforms.ToTensor())

# 加载验证集
MNIST_val = torchvision.datasets.MNIST(data_path, train=False, download=True,
                                              transform=torchvision.transforms.ToTensor())

# 创建数据集迭代器，用于实现批量处理数据
train_loader = torch.utils.data.DataLoader(MNIST, batch_size=100, shuffle=True)
val_loader = torch.utils.data.DataLoader(MNIST_val, batch_size=100, shuffle=True)

# 定义LeNet模型
class LeNet(nn.Module):
    def __init__(self, intput_channle=1, output_channle=10):
        super().__init__()
        self.conv1 = nn.Conv2d(intput_channle, 6, kernel_size=5, padding=2)
        self.act1 = nn.Sigmoid()
        self.pool1 = nn.AvgPool2d(2)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.act2 = nn.Sigmoid()
        self.pool2 = nn.AvgPool2d(2)
        self.flatten = nn.Flatten()
        self.linear1 = nn.Linear(16 * 5 * 5, 120)
        self.act3 = nn.Sigmoid()
        self.linear2 = nn.Linear(120, 84)
        self.act4 = nn.Sigmoid()
        self.linear3 = nn.Linear(84, output_channle)

    def forward(self, x):
        out = self.pool1(self.act1(self.conv1(x)))
        out = self.pool2(self.act2(self.conv2(out)))
        out = self.flatten(out)
        out = self.act3(self.linear1(out))
        out = self.act4(self.linear2(out))
        out = self.linear3(out)
        return out

# 定义循环训练函数
def train(n_epochs, optimizer, model, loss_fn, train_loader):
    """
    定义模型训练函数，实现前向运算与后向传播训练
    Args:
        n_epochs (_type_): 训练的总周期数
        optimizer (_type_): 训练优化器
        model (_type_): 模型
        loss_fn (_type_): 损失函数
        train_loader (_type_): 数据迭代器
    """
    for epoch in range(1, n_epochs+1):
        for imgs, labels in tqdm(train_loader):
            # 将imgs与labels放到所选设备上
            imgs = imgs.to(device=device)
            labels = labels.to(device=device)
            # 执行前向运算
            outputs = model(imgs)
			# 使用预测值outputs与实际值labels计算loss
            loss = loss_fn(outputs, labels)
			
            # 注意在loss.backward()之前一定要将优化器的梯度置零，
            # 因为pytorch训练时优化器的梯度是累计的，在一批量迭代结束后，
            # 梯度并没有归零，如果不置零，则梯度会一直累加。
            optimizer.zero_grad()
            # 
            loss.backward()
            optimizer.step()
        print("{} 第 {} epoch的loss值为：{}".format(datetime.datetime.now(), epoch, loss))


def validata(model, train_loader, val_loader):
    """
    定义验证函数（本次既验证了训练集，又验证了验证集）

    Args:
        model (_type_): 模型
        train_loader (_type_): 训练迭代器
        val_loader (_type_): 验证迭代器
    """
    for name, loader in [("train", train_loader), ("test", val_loader)]:
        correct = 0
        total = 0
        # with torch.no_grad():下面的内容都不会涉及梯度更新（验证时只执行前向传播，不需要进行梯度更新）。
        with torch.no_grad():
            for imgs, labels in tqdm(loader):
                imgs = imgs.to(device=device)
                labels = labels.to(device=device)
                # 执行前向传播，得到预测结果
                outputs = model(imgs)
                # 获取预测结果的最大值的索引（返回的第一个tensor是最大值，第二个tensor是最大值的索引）
                _, predicted = torch.max(outputs, dim=1)
				# 计算模型预测的样本总数
                total += labels.shape[0]
                # 迭代计算模型的预测正确的综述
                correct += int((predicted == labels).sum())
        # 计算模型预测的精度（此处是准确率）
        print("Accuracy {} : {:.5f}".format(name, correct / total))

# 创建模型实例，并且将模型放在所选设备上
LeNet = LeNet().to(device=device)
# 加载预训练权重（本次加载的是我之前训练保存的权重）
LeNet.load_state_dict(torch.load(pt_path + 'ptMNIST.pt'))

# 创建优化器（本次使用的是Adam优化器，学习率设置为0.01）
optimiter = torch.optim.Adam(LeNet.parameters(), lr=0.01)

# 创建loss函数（本次使用是交叉熵损失函数，参数全保持默认）
loss_fn = nn.CrossEntropyLoss()

# 开始训练，训练5个epoch
train(n_epochs=5,
      optimizer=optimiter,
      model=LeNet,
      loss_fn=loss_fn,
      train_loader=train_loader)
# 保存模型训练权重
torch.save(LeNet.state_dict(), pt_path + 'MNIST.pt')

# 开始验证
validata(model=LeNet, train_loader=train_loader, val_loader=val_loader)

训练结果展示如下：

注意：本次所展示的流程仅是最基本的神经网络训练实现流程，很多细节的地方并没有展开细说，只是希望大家能够通过本节内容，了解模型训练验证的大概流程，心里有个大体框架，在之后的学习中能有个好的起点。

在以后的内容中，我们会对深度学习（尤其是计算机视觉方向）展开更深入的学习介绍，更多的模型架构，例如ConNext、 Vison Transformer等等、更多的优化器、学习率调整策略、loss函数、模型可视化、模型训练tricks、模型部署方法等内容，会从很多细节出发，进行讲解介绍。以及会带来更多的实战内容（分类、分割、检测等等）。我会和大家一起更深入广泛地学习深度学习！