Pytorch训练LeNet模型MNIST数据集

如何用torch框架训练深度学习模型(详解)

0. 需要的包

python 复制代码
import torch
from torch.nn import CrossEntropyLoss
from torch.optim import SGD
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

1. 数据加载和导入

以MNIST数据集为例

python 复制代码
# 1.1 需要设置数据归一化
train_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))])
test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))])
# 1.2 用dataset.MNIST函数下载和加载训练集与测试集 
train_dataset = datasets.MNIST(dataset_path, train=True, 
	download=False, transform=train_transform)
test_dataset = datasets.MNIST(dataset_path, train=False, 
	download=False, transform=test_transform)
# 1.3 加载进dataload用于后续数据按batch取用
batch_size = 256
train_loader = DataLoader(train_dataset, batch_size=batch_size)
test_loader = DataLoader(test_dataset, batch_size=batch_size)

补充:这里的transform根据不同的数据集选择不同的值

datasets加载数据集时path的路径为:'.\data\' 该目录下包括\MNIST文件夹

2. 加载模型和设置超参数

python 复制代码
# 2.1 这里需要提前定义model的class,包括层结构和forward函数
model = LeNet_Mnist().to(device)
# 2.2 设置优化器、损失函数、训练轮次
learning_rate = 1e-2
# 传入模型参数,用于优化更新
sgd = SGD(model.parameters(), lr=learning_rate)  
loss_fn = CrossEntropyLoss()
all_epoch = 20

3. 训练

python 复制代码
# 3.1 首先设置训练模式
model.train()
# 3.2 按照batch从train_loader中批量选择数据
for idx, (train_x, train_label) in enumerate(train_loader):
    train_x = train_x.to(device)
    train_label = train_label.to(device)
    sgd.zero_grad()
    predict_y = model(train_x.float())
    loss = loss_fn(predict_y, train_label.long())
    loss.backward()
    sgd.step()

补充:可以在外面再套一层迭代次数

python 复制代码
for current_epoch in range(all_epoch):  # local training

4. 测试

python 复制代码
# 4.1 记录测试结果
all_correct_num = 0
all_sample_num = 0
# 4.2 进入模型验证模式,该模式下不会修改梯度
model.eval()
# 4.3 按批次测试
for idx, (test_x, test_label) in enumerate(test_loader):
    test_x = test_x.to(device)
    test_label = test_label.to(device)
    predict_y = model(test_x.float()).detach()
    predict_y = torch.argmax(predict_y, dim=-1)
    current_correct_num = predict_y == test_label
    all_correct_num += np.sum(current_correct_num.to('cpu').numpy(), axis=-1)
    all_sample_num += current_correct_num.shape[0]
# 4.4 记录结果并输出
acc = all_correct_num / all_sample_num
print('accuracy: {:.3f}'.format(acc), flush=True)

5. 保存结果

python 复制代码
# 5.1 保存参数
print("Save the model state dict")
torch.save(model.state_dict(), "./lenet_mnist.pt")
# 5.2 或者也可以选择保存checkpoint,每轮都保存一次,万一中断能继续
checkpoint = {
                "model": model.state_dict(),
                "optim": sgd.state_dict(),
             }
print("Save the checkpoint")
torch.save(checkpoint, "./checkpoint{}.pt".format(current_epoch))
相关推荐
机器之心11 分钟前
高阶程序,让AI从技术可行到商业可信的最后一公里
人工智能·openai
martinzh15 分钟前
解锁RAG高阶密码:自适应、多模态、个性化技术深度剖析
人工智能
机器之心19 分钟前
刚刚,李飞飞空间智能新成果震撼问世!3D世界生成进入「无限探索」时代
人工智能·openai
scilwb29 分钟前
Isaac Sim机械臂教程 - 阶段1:基础环境搭建与机械臂加载
人工智能·开源
舒一笑1 小时前
TorchV企业级AI知识引擎的三大功能支柱:从构建到运营的技术解析
人工智能
掘金酱1 小时前
🎉 2025年8月金石计划开奖公示
前端·人工智能·后端
鹏多多2 小时前
纯前端人脸识别利器:face-api.js手把手深入解析教学
前端·javascript·人工智能
aneasystone本尊2 小时前
盘点 Chat2Graph 中的专家和工具
人工智能
这里有鱼汤2 小时前
小白必看:QMT里的miniQMT入门教程
后端·python
Baihai_IDP3 小时前
AI Agents 能自己开发工具自己使用吗?一项智能体自迭代能力研究
人工智能·面试·llm