Pytorch训练LeNet模型MNIST数据集

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

0. 需要的包

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数据集为例

# 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. 加载模型和设置超参数

# 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. 训练

# 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()

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

for current_epoch in range(all_epoch):  # local training

4. 测试

# 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. 保存结果

# 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))

6. 导入模型

model.load_state_dict(torch.load(resume_path))

resume_path为保存模型的地址