PyTorch MNIST手写体识别：SwanLab可视化训练

MNIST手写体识别是深度学习、CV的“Hello World”，无数人从这个任务入门，进入深度学习的世界～

进阶请看：PyTorch猫狗分类

这篇文章我将带大家使用SwanLab（训练可视化）、PyTorch（深度学习框架）这两个开源工具，完成从数据集准备、代码编写、到模型训练的全过程。 [图片]

代码同时适用于纯CPU、英伟达GPU或Apple M1训练。

完整代码请直接跳转到2.9节。

实验过程：SwanLab - MNIST手写体识别数据集：百度云提取码: 89y2 两个开源库：SwanLab、pytorch

1. 准备部分

在此之前，你需要确保你已经安装了Python。

1.1 什么是MNIST手写体识别

MNIST手写体识别任务是一个经典的计算机视觉问题，属于图像分类任务，目标是输入一个手写数字图像，AI模型可以正确预测数字是多少。 MNIST数据集包含70,000个手写数字图像，每个图像大小为28x28像素。这些图像分为两部分：60,000个训练集和10,000个测试集。 [图片]

1.2 安装Python库

需要安装下面这4个库：

shell复制代码torch
torchvision
swanlab>=0.3.2

安装命令：

复制代码pip install torch torchvision swanlab

1.3 创建文件目录

需要包含两个文件：

train.py，它的作用训练模型
MNIST文件夹，存放数据集。

MNIST数据集的下载有两种方式：一种是通过2.1中的代码下载（推荐），如果下载存在网络问题的话，则可以用我提供的百度云链接下载：百度云，提取码: 89y2。

[图片]

2. 训练部分

如果想直接看完整代码和效果，可直接跳转到第2.9。

2.1 载入MNIST数据集

训练部分的代码全部写在train.py中。载入MNIST数据集分为三步：

下载MNIST数据集
将数据集分为训练集和验证集
封装到PyTorch DataLoader中，以供训练使用

下面的代码演示了这一过程：

python复制代码from torchvision.datasets import MNIST

# 下载MNIST训练集
dataset = MNIST(os.getcwd(), train=True, download=True, transform=ToTensor())

# 随机分为训练集和验证集, 训练集55000张，验证集5000张
train_dataset, val_dataset = utils.data.random_split(dataset, [55000, 5000])

# 写入Pytorch数据载入器中
train_loader = utils.data.DataLoader(train_dataset, batch_size=256, shuffle=True)
val_loader = utils.data.DataLoader(val_dataset, batch_size=1, shuffle=False)

ps：如果下载速度不给力的话，可以用我提供的百度云链接下载：百度云提取码: 89y2。

目录结构是 根目录/MNIST/raw/...(下图所示的文件)。在这里插入图片描述

2.2 载入ResNet18模型

模型我们直接选用经典的ResNet18，模型的具体原理本文不细说，重点放在工程实现上。我们使用torchvision来创建1个resnet50模型，并载入预训练权重：

python复制代码import torchvision
from torchvision.models import ResNet18_Weights

model = torchvision.models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

因为MNIST是个1维图像和10分类任务（即输入一张图片，分类它是0～9这10个数字中的哪一个），而torchvision提供的resnet18默认是适配3维图像与1000分类，所以我们需要把模型的第一个卷积层的输入维度替换为1，全连接层的输出维度替换为10：\

python复制代码import torchvision
from torchvision.models import ResNet18_Weights
import torch

model = torchvision.models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

# 将模型的输入维度改为1，全连接层替换为10
model.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
in_features = model.fc.in_features
model.fc = torch.nn.Linear(in_features, 10)

2.3 设置cuda/mps/cpu

如果你的电脑是英伟达显卡，那么cuda可以极大加速你的训练；如果你的电脑是Macbook Apple Sillicon（M系列芯片），那么mps同样可以极大加速你的训练；如果都不是，那就选用cpu：

python复制代码#检测是否支持mps
try:
    use_mps = torch.backends.mps.is_available()
except AttributeError:
    use_mps = False

#检测是否支持cuda
if torch.cuda.is_available():
    device = "cuda"
elif use_mps:
    device = "mps"
else:
    device = "cpu"

将模型加载到对应的device中：

python

复制代码model.to(torch.device(device))

2.4 初始化SwanLab

在训练中我们使用swanlab库作为实验管理与指标可视化工具。 SwanLab是一个类似Wandb（Weights and Biases）的在线开源训练可视化工具，能够在远程、在网页上看指标。除了能记录指标，还能自动记录训练的logging、硬件环境、Python环境、训练时间等信息。如何入门SwanLab可以看这个文档：🚀快速开始 - SwanLab，也可以跟着本文往下走。

初始化并设置超参数 swanlab库使用swanlab.init记录超参数、实验名、实验介绍：

python复制代码import swanlab
# 初始化swanlab
run = swanlab.init(
    project="MNIST-example",
    experiment_name="ResNet18",
    config={
        "model": "ResNet18",
        "optim": "Adam",
        "lr": 1e-4,
        "batch_size": 256,
        "num_epochs": 10,
        "device": device,
    }

这里我们设置的超参数分别是：模型使用resnet18，学习率lr为1e-4，批大小batch_size为256，轮数epoch为20，分类数为10，设备为2.3检测到的设备值。这里swanlab干了三件事：

创建了1个名为MNIST的项目
创建了1个名为ResNet18的实验，项目和实验是类似文件夹和文件的关系，每次训练都会产生一个新的实验
将超参数上传到实验中，被记录下来

跟踪训练指标 在训练MNIST的过程中，我们最关心的指标就是训练集的损失值loss和验证集的准确率acc，我们用SwanLab在记录这些指标，生成可视化的折线图。下面是一个简单的记录loss的常见用法，具体的使用案例见2.6和2.7。

python复制代码for data in train_dataloader:
    ...
    swanlab.log({"loss": loss})

2.5 设置优化器和损失函数

设置损失函数为交叉熵损失，优化器为Adam。

python复制代码criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=run.config.lr)

交叉熵损失是常用于图像分类任务的损失函数，Adam是一种经典的优化器。

2.6 定义训练函数train

我们定义1个训练函数train：

python复制代码def train(model, device, train_dataloader, optimizer, criterion, epoch, num_epochs):
    model.train()
    # 1. 循环调用train_dataloader，每次取出1个batch_size的图像和标签
    for iter, (inputs, labels) in enumerate(train_dataloader):
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        # 2. 传入到resnet18模型中得到预测结果
        outputs = model(inputs)
        # 3. 将结果和标签传入损失函数中计算交叉熵损失
        loss = criterion(outputs, labels)
        # 4. 根据损失计算反向传播
        loss.backward()
        # 5. 优化器执行模型参数更新
        optimizer.step()
        print('Epoch [{}/{}], Iteration [{}/{}], Loss: {:.4f}'.format(epoch, num_epochs, iter + 1, len(train_dataloader),
                                                                      loss.item()))
        # 6. 每20次迭代，用SwanLab记录一下loss的变化
        if iter % 20 == 0:
            swanlab.log({"train/loss": loss.item()})

训练的逻辑很简单：我们循环调用train_dataloader，每次取出1个batch_size的图像和标签，传入到resnet18模型中得到预测结果，将结果和标签传入损失函数中计算交叉熵损失，最后根据损失计算反向传播，Adam优化器执行模型参数更新，循环往复。在训练中我们最关心的指标是损失值loss，所以我们用swanlab.log跟踪它的变化。

2.7 定义测试函数test

我们定义1个测试函数test：

python复制代码def test(model, device, val_dataloader, epoch):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        # 1. 循环调用val_dataloader，每次取出1个batch_size的图像和标签
        for inputs, labels in val_dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            # 2. 传入到resnet18模型中得到预测结果
            outputs = model(inputs)
            # 3. 获得预测的数字
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            # 4. 计算与标签一致的预测结果的数量
            correct += (predicted == labels).sum().item()
    
        # 5. 得到最终的测试准确率
        accuracy = correct / total
        # 6. 用SwanLab记录一下准确率的变化
        swanlab.log({"val/accuracy": accuracy}, step=epoch)

测试的逻辑同样很简单：我们循环调用test_dataloader，将验证集的图像传入到当前训练的resnet18模型中得到预测结果，与标签进行对比，计算整体的准确率。在测试中我们最关心的指标是准确率accuracy，所以我们用swanlab.log跟踪它的变化。

2.8 保存模型权重文件

我们使用train.save进行模型的权重保存，保存到checkpoint文件夹下。

python复制代码# 如果不存在checkpoint文件夹，则自动创建一个
if not os.path.exists("checkpoint"):
    os.makedirs("checkpoint")
torch.save(model.state_dict(), 'checkpoint/latest_checkpoint.pth')

至此，我们完成已完成了绝大多数的代码，现在是时候将它们组合起来，开始训练！

2.9 完整训练代码

我们一共训练10轮，每2轮进行测试，并在最后保存权重文件：

python复制代码import os
import torch
from torch import nn, optim, utils
import torch.nn.functional as F
import torchvision
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torchvision.models import ResNet18_Weights
import swanlab

def train(model, device, train_dataloader, optimizer, criterion, epoch, num_epochs):
    model.train()
    # 1. 循环调用train_dataloader，每次取出1个batch_size的图像和标签
    for iter, (inputs, labels) in enumerate(train_dataloader):
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        # 2. 传入到resnet18模型中得到预测结果
        outputs = model(inputs)
        # 3. 将结果和标签传入损失函数中计算交叉熵损失
        loss = criterion(outputs, labels)
        # 4. 根据损失计算反向传播
        loss.backward()
        # 5. 优化器执行模型参数更新
        optimizer.step()
        print('Epoch [{}/{}], Iteration [{}/{}], Loss: {:.4f}'.format(epoch, num_epochs, iter + 1, len(train_dataloader),
                                                                      loss.item()))
        # 6. 每20次迭代，用SwanLab记录一下loss的变化
        if iter % 20 == 0:
            swanlab.log({"train/loss": loss.item()})

def test(model, device, val_dataloader, epoch):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        # 1. 循环调用val_dataloader，每次取出1个batch_size的图像和标签
        for inputs, labels in val_dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            # 2. 传入到resnet18模型中得到预测结果
            outputs = model(inputs)
            # 3. 获得预测的数字
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            # 4. 计算与标签一致的预测结果的数量
            correct += (predicted == labels).sum().item()
    
        # 5. 得到最终的测试准确率
        accuracy = correct / total
        # 6. 用SwanLab记录一下准确率的变化
        swanlab.log({"val/accuracy": accuracy}, step=epoch)
    

if __name__ == "__main__":

    #检测是否支持mps
    try:
        use_mps = torch.backends.mps.is_available()
    except AttributeError:
        use_mps = False

    #检测是否支持cuda
    if torch.cuda.is_available():
        device = "cuda"
    elif use_mps:
        device = "mps"
    else:
        device = "cpu"

    # 初始化swanlab
    run = swanlab.init(
        project="MNIST-example",
        experiment_name="ResNet18",
        config={
            "model": "ResNet18",
            "optim": "Adam",
            "lr": 1e-4,
            "batch_size": 256,
            "num_epochs": 10,
            "device": device,
        },
    )

    # 设置MNIST训练集和验证集
    dataset = MNIST(os.getcwd(), train=True, download=True, transform=ToTensor())
    train_dataset, val_dataset = utils.data.random_split(dataset, [55000, 5000])

    train_dataloader = utils.data.DataLoader(train_dataset, batch_size=run.config.batch_size, shuffle=True)
    val_dataloader = utils.data.DataLoader(val_dataset, batch_size=8, shuffle=False)

    # 初始化模型
    model =torchvision.models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
    # 让模型适配MNIST数据集
    model.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
    in_features = model.fc.in_features
    model.fc = torch.nn.Linear(in_features, 10)
    model.to(torch.device(device))

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=run.config.lr)

    # 开始训练和测试循环
    for epoch in range(1, run.config.num_epochs+1):
        swanlab.log({"train/epoch": epoch}, step=epoch)
        train(model, device, train_dataloader, optimizer, criterion, epoch, run.config.num_epochs)
        if epoch % 2 == 0: 
            test(model, device, val_dataloader, epoch)

    # 保存模型
    # 如果不存在checkpoint文件夹，则自动创建一个
    if not os.path.exists("checkpoint"):
        os.makedirs("checkpoint")
    torch.save(model.state_dict(), 'checkpoint/latest_checkpoint.pth')