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个猫狗分类模型的全部过程,更多想了解的可以参考相关链接或评论此文章。 如果有帮助,请点个赞和收藏吧~
最后的最后
感谢你们的阅读和喜欢,我收藏了很多技术干货,可以共享给喜欢我文章的朋友们,如果你肯花时间沉下心去学习,它们一定能帮到你。
因为这个行业不同于其他行业,知识体系实在是过于庞大,知识更新也非常快。作为一个普通人,无法全部学完,所以我们在提升技术的时候,首先需要明确一个目标,然后制定好完整的计划,同时找到好的学习方法,这样才能更快的提升自己。
这份完整版的大模型 AI 学习资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费
】
大模型知识脑图
为了成为更好的 AI大模型 开发者,这里为大家提供了总的路线图。它的用处就在于,你可以按照上面的知识点去找对应的学习资源,保证自己学得较为全面。
经典书籍阅读
阅读AI大模型经典书籍可以帮助读者提高技术水平,开拓视野,掌握核心技术,提高解决问题的能力,同时也可以借鉴他人的经验。对于想要深入学习AI大模型开发的读者来说,阅读经典书籍是非常有必要的。
实战案例
光学理论是没用的,要学会跟着一起敲,要动手实操,才能将自己的所学运用到实际当中去,这时候可以搞点实战案例来学习。
面试资料
我们学习AI大模型必然是想找到高薪的工作,下面这些面试题都是总结当前最新、最热、最高频的面试题,并且每道题都有详细的答案,面试前刷完这套面试题资料,小小offer,不在话下
640套AI大模型报告合集
这套包含640份报告的合集,涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师,还是对AI大模型感兴趣的爱好者,这套报告合集都将为您提供宝贵的信息和启示。
这份完整版的大模型 AI 学习资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费
】