一、标准训练流程#

1. 数据准备与长度获取#

在开始训练前，需明确数据集规模，以便在日志中计算进度。

1
# 准备数据集
2
train_data = torchvision.datasets.CIFAR10(root="../data", train=True, transform=torchvision.transforms.ToTensor(), download=True)
3
test_data = torchvision.datasets.CIFAR10(root="../data", train=False, transform=torchvision.transforms.ToTensor(), download=True)
4

5
# 获取长度：用于后续计算平均 Loss 或准确率
6
train_data_size = len(train_data)
7
test_data_size = len(test_data)
8
print(f"训练集长度：{train_data_size}, 测试集长度：{test_data_size}")
9

10
# 加载数据
11
train_dataloader = DataLoader(train_data, batch_size=64)
12
test_dataloader = DataLoader(test_data, batch_size=64)

2. 训练配置与日志监控 (TensorBoard)#

使用 SummaryWriter 记录训练曲线是调试模型的必备手段。

1
# 损失函数与优化器
2
loss_fn = nn.CrossEntropyLoss()
3
learning_rate = 1e-2 # 即 0.01
4
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
5

6
# 设置辅助参数
7
total_train_step = 0 # 记录总训练步数
8
total_test_step = 0  # 记录测试步数
9
epoch = 10           # 训练轮数
10

11
# 实例化 TensorBoard
12
writer = SummaryWriter("../logs_train")

二、核心训练与测试套路#

1. 训练阶段 (Training)#

在每一轮 Epoch 中，通过计数器实现定时日志输出，避免控制台刷新过快。

1
for i in range(epoch):
2
    print(f"-------第 {i+1} 轮训练开始-------")
3
    model.train()
4
    for data in train_dataloader:
5
        imgs, targets = data
6
        outputs = model(imgs)
7
        loss = loss_fn(outputs, targets)
8

9
        # 优化器优化
10
        optimizer.zero_grad()
11
        loss.backward()
12
        optimizer.step()
13

14
        total_train_step += 1
15
        if total_train_step % 100 == 0:
16
            print(f"训练次数：{total_train_step}, Loss: {loss.item()}")
17
            writer.add_scalar("train_loss", loss.item(), total_train_step)

2. 测试与准确率计算 (Evaluation)#

测试集的表现是评估模型泛化能力的关键。准确率计算逻辑为：(预测正确的样本数) / (总样本数)。

1
    model.eval()
2
    total_test_loss = 0
3
    total_accuracy = 0
4
    with torch.no_grad():
5
        for data in test_dataloader:
6
            imgs, targets = data
7
            outputs = model(imgs)
8
            loss = loss_fn(outputs, targets)
9
            total_test_loss += loss.item()
10

11
            # 计算分类正确的数量
12
            accuracy = (outputs.argmax(1) == targets).sum()
13
            total_accuracy += accuracy
14

15
    print(f"测试集总 Loss: {total_test_loss}")
16
    print(f"测试集准确率: {total_accuracy / test_data_size}")
17

18
    # 记录到 TensorBoard
19
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
20
    writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)
21
    total_test_step += 1

三、模型持久化与 GPU 加速#

1. 循环保存策略#

在每一轮 Epoch 结束后保存模型，防止意外中断导致训练成果丢失。

1
    # 方式一：保存整个模型
2
    torch.save(model, f"model_{i}.pth")
3
    print("模型已保存")

2. GPU 硬件加速#

通过 .to(device) 确保所有计算单元位于同一设备。

1
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
2

3
model = model.to(device)
4
loss_fn = loss_fn.to(device)
5

6
# 数据迁移
7
for data in train_dataloader:
8
    imgs, targets = data
9
    imgs, targets = imgs.to(device), targets.to(device)