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')
2.10 开始训练!
实验过程可看这个SwanLab网页:SwanLab - MNIST 在开始运行时,如果你是第一次使用SwanLab,会提示你要输入SwanLab API Key。 这里的操作也非常简单,去SwanLab官网登录一下账号,在【设置】页面复制API Key,粘贴过来就可以:
然后,我们运行train.py:
这时候你会在看到在开头会给到你两个链接,我们点击第一个,里面包含了这个项目的信息和一个对比实验表格:
我们点开1个进行中的实验,会看到loss和acc整体的变化曲线:
切换到实验卡片,这里记录了实验的各种信息,包括超参数、最终的实验指标、实验状态、训练时长、Git仓库链接、主机名、操作系统、Python版本、硬件配置等等。
可以看到我们最终训练的准确率是98.66%,提高到99%以上应该不困难!可以把这个实验作为你的baseline,去尝试不同的trick来提高准确率~
至此我们完成了模型的训练和测试,得到了1个表现非常棒的猫狗分类模型,权重保存在了checkpoint目录下。 至此,我们完成了用SwanLab、PyTorch这两个开源工具训练1个猫狗分类模型的全部过程,更多想了解的可以参考相关链接或评论此文章。 如果有帮助,请点个赞和收藏吧~
读者福利:如果大家对大模型感兴趣,这套大模型学习资料一定对你有用
对于0基础小白入门:
如果你是零基础小白,想快速入门大模型是可以考虑的。
一方面是学习时间相对较短,学习内容更全面更集中。
二方面是可以根据这些资料规划好学习计划和方向。
资源分享
大模型AGI学习包
资料目录
- 成长路线图&学习规划
- 配套视频教程
- 实战LLM
- 人工智能比赛资料
- AI人工智能必读书单
- 面试题合集
《人工智能\大模型入门学习大礼包》,可以扫描下方二维码免费领取!
![](https://img-blog.csdnimg.cn/img_convert/d2d2c00b423aceb279f3955eb0b5386e.png)
1.成长路线图&学习规划
要学习一门新的技术,作为新手一定要先学习成长路线图,方向不对,努力白费。
对于从来没有接触过网络安全的同学,我们帮你准备了详细的学习成长路线图&学习规划。可以说是最科学最系统的学习路线,大家跟着这个大的方向学习准没问题。
2.视频教程
很多朋友都不喜欢晦涩的文字,我也为大家准备了视频教程,其中一共有21个章节,每个章节都是当前板块的精华浓缩。
3.LLM
大家最喜欢也是最关心的LLM(大语言模型)
《人工智能\大模型入门学习大礼包》,可以扫描下方二维码免费领取!
![](https://img-blog.csdnimg.cn/img_convert/d2d2c00b423aceb279f3955eb0b5386e.png)