如何用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_labelall_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))