适合入门深度学习的项目！！！搭建LeNet-5网络模型——2 训练模型搭建

上一节里，我们的网络模型放在一个名为net.py的文件里，这里我们调用它进行了模型的训练和测试，并将最优的模型保存在指定路径里。

详细注释的代码如下：

import torch
from torch import nn
from net import MyLeNet
from torch.optim import lr_scheduler  # 学习率经过多次会下降，使得学习率会变得更小
from torchvision import datasets, transforms
import os

# # 数据转换为tensor格式。数据刚拿到是矩阵(array)的格式，但神经网络里面需要张量(tensor)的格式
# data_transform = transforms.Compose(
#     [transforms.ToTensor]
# )
data_transform = transforms.ToTensor()

# 加载预训练的数据集
train_dataset = datasets.MNIST(root='./data',train=True,transform=data_transform,download=True)
train_dataloader = torch.utils.data.DataLoader(dataset = train_dataset,batch_size = 16, shuffle = True)

# 加载测集的数据
test_dataset = datasets.MNIST(root='./data',train=False,transform=data_transform,download=True)
test_dataloader = torch.utils.data.DataLoader(dataset = test_dataset,batch_size = 16, shuffle = True)

# 选择设备
# Get cpu, gpu or mps device for training.
device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")

# 调用net(1 网络模型的文件)里面定义的模型，将模型数据转到GPU
model = MyLeNet().to(device)

# 定义一个损失函数(交叉熵)
loss_fn = nn.CrossEntropyLoss()

# 定义一个优化器（随机梯度下降）

# SGD(需要反向传播的参数，学习率,momentum动量)
# momentum 动量是依据物理学的势能与动能之间能量转换原理提出来的。
# 当 momentum 动量越大时，其转换为势能的能量也就越大，就越有可能摆脱局部凹域的束缚，进入全局凹域。
# 大概意思就是：在随机梯度下降的过程中，一开始的幅度会很大，如果加上该动量
# 梯度下降的方向上会缓慢下降，容易找到我们需要的最低点，更好收敛
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)

# 当学习率（Learning Rate,lr）设置的过小时，收敛过程将变得十分缓慢;
# 而当lr设置的过大时，梯度可能会在最小值附近来回震荡，甚至可能无法收敛。
# 合理的lr调整可以在梯度下降过程防止抖动比较大，越往后面学，学习率越小越容易找到最小点，越好收敛
# 学习率调整有离散下降、指数减缓、分数减缓这三种常用方法。参考：http://t.csdn.cn/UkBFJ

# 这里我们采用离散下降：每过10轮，lr变为原来的0.1
lr_scheduler = lr_scheduler.StepLR(optimizer,step_size=10,gamma=0.1)

# 定义训练函数
def train(dataloader, model, loss_fn, optimizer):
    loss, correct, n = 0.0, 0.0, 0
    # 批次，数据(eg:图片，标签)
    for batch,(X,y) in enumerate(dataloader):
        # 1.前向传播

        # 先传入指定设备中
        X,y = X.to(device),y.to(device)
        output = model(X)

        # 损失函数——交叉熵，(输出值，真实的lable)，用于反向传播
        cur_loss = loss_fn(output, y)

        # 用于后面计算loss值
        _, pred = torch.max(output, axis = 1)

        # 一组数据有16个批次——output.shape[0]
        # 其中准确的批次数——torch.sum(y == pred)
        cur_acc = torch.sum(y == pred)/output.shape[0]

        # 2.反向传播
        optimizer.zero_grad() # 梯度清零
        cur_loss.backward() # loss值越大反向传播力度越大，loss值越小说明效果越好
        optimizer.step() #梯度更新

        # 计算每一batch的loss
        # item():取出单元素张量的元素值并返回该值，保持原元素类型不变
        loss += cur_loss.item()
        correct +=  cur_acc.item()
        # 统计n,计算平均精确度
        n = n + 1
    loss /= n
    correct /= n
    print(f"Test Error: \n Train Accuracy: {(100 * correct):>0.1f}%, Train Avg_loss: {loss:>8f} \n")

# 定义测试函数
def val(dataloader, model, loss_fn):
    model.eval() # 验证模式
    loss, correct, n = 0.0, 0.0, 0
    # 模型不参与更新，用with torch.no_grad()
    # 验证集的时候，我们只是想看一下训练的效果，并不是想通过验证集来更新网络时
    # 使用该语句不会影响之前的量
    with torch.no_grad():
        for batch,(X,y) in enumerate(dataloader):
            # 1.前向传播（和train内容相似）
            X,y = X.to(device),y.to(device)
            output = model(X)
            cur_loss = loss_fn(output, y)
            _, pred = torch.max(output, axis = 1)
            cur_acc = torch.sum(y == pred)/output.shape[0]
            # 2.测试里没有反向传播去更新参数奥
            loss += cur_loss.item()
            correct +=  cur_acc.item()
            # 统计n,计算平均精确度
            n = n + 1
        loss /= n
        correct /= n
        print(f"Test Error: \n Val Accuracy: {(100 * correct):>0.1f}%, Val Avg_loss: {loss:>8f} \n")
        # 返回精确度，用以保存指定效果精确度的模型
        return correct
# 训练
epochs = 50
max_acc = 0
for t in range(epochs):
    print(f"Epoch {t + 1}\n-------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    a = val(test_dataloader, model, loss_fn)
    # 保存最好的模型权重
    if a > max_acc:
        folder = 'save_model'
        if not os.path.exists(folder):
            os.mkdir('save_model')
        min_acc = a
        print('save best model')
        torch.save(model.state_dict(),'save_model/best.pth')
print("Done!")

经过50个epoch的训练，模型达到了98.2%的准确率：

 参考：B站视频 【从0开始撸代码--手把手教你搭建LeNet-5网络模型】