Pytorch深入浅出（十四）之完整的模型训练测试套路

深度学习的"五步走"流程（数据 -> 模型 -> 损失函数 -> 优化器 -> 迭代训练 ）

迭代训练（Step 5）有很多"套路"和行业标准写法，现在我们由浅入深的进行讲解。

迭代训练的几大要点

"数据 → 模型 → 损失函数 → 优化器 → 迭代训练"是算法层面的闭环 ，但工程上，"迭代训练"会被展开成：

epoch 级循环
模式切换（train / eval）
梯度控制（zero / backward / step）
指标统计
日志 & 保存

CIFAR10数据集分类

在该示例中，我们使用 CIFAR10 数据集完成一个经典的图像分类任务。

该示例的目的不是追求复杂的数据工程，而是帮助理解深度学习"五步走"中 Step5：迭代训练的基本结构与语义。

python 复制代码

from torch.utils.tensorboard import SummaryWriter
import torchvision
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import models

# ==================== Step 1：数据 ====================
train_data = torchvision.datasets.CIFAR10(
    "../dataset",
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=True
)

test_data = torchvision.datasets.CIFAR10(
    "../dataset",
    train=False,
    transform=torchvision.transforms.ToTensor(),
    download=True
)

train_data_size = len(train_data)
test_data_size = len(test_data)
print(f"训练数据集的长度为：{train_data_size}")
print(f"测试数据集的长度为：{test_data_size}")

train_dataloader = DataLoader(train_data, batch_size=64, shuffle=True, drop_last=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=False, drop_last=True)

# ==================== Step 2：模型 ====================
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 10)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# ==================== Step 3：损失函数 ====================
loss_fn = nn.CrossEntropyLoss()
# loss_fn.to(device)

# ==================== Step 4：优化器 ====================
learning_rate = 1e-2
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

# ==================== Step 5：迭代训练 ====================
total_train_step = 0
total_test_step = 0
epoch = 10

writer = SummaryWriter("../logs_train")

for i in range(epoch):
    print(f"-------------第 {i + 1} 轮训练开始------------")

    # ========= 训练阶段 =========
    model.train()
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)

        outputs = model(imgs)
        loss = loss_fn(outputs, targets)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

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

    # ========= 测试阶段 =========
    model.eval()
    total_test_loss = 0.0

    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)

            outputs = model(imgs)
            loss = loss_fn(outputs, targets)
            total_test_loss += loss.item()

    avg_test_loss = total_test_loss / len(test_dataloader)
    print(f"整体测试集上的平均 Loss：{avg_test_loss:.4f}")
    writer.add_scalar("test_loss", avg_test_loss, total_test_step)
    total_test_step += 1

    torch.save(model.state_dict(), f"model_{i}.pth")
    print("模型已保存")

writer.close()

与后面的树叶分类竞赛不同，CIFAR10 是一个标准、干净、已标注好的数据集，因此：

对"Step1：数据"这步，数据处理非常简单：1.使用 torchvision.datasets.CIFAR10 直接加载数据集；2.使用 transforms.ToTensor() 将图片转为 Tensor；3.使用 DataLoader进行迭代

对"Step5：迭代训练"这步，训练和测试都放在同一个 epoch 内完成，实现了最小但完整的训练闭环，同时：统计训练步数；每 100 step 记录一次 loss；使用 TensorBoard 进行日志记录。每个 epoch 保存一次模型（教学级写法），在更规范的工程/竞赛中，通常只保存验证集表现最好的模型。

树叶分类竞赛

在该示例中，我们使用Kaggle竞赛的树叶分类数据集完成一个图像分类任务。

该示例已经是一个完整的工程级复现，以及地带训练模板可复用。

python 复制代码

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader, random_split
import pandas as pd
from PIL import Image
from sklearn.preprocessing import LabelEncoder
from torchvision import transforms, models
import os
from tqdm import tqdm

# 定义设备（优先使用GPU，如果没有则使用CPU）
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {DEVICE}")

# ==================== 配置参数 ====================
ROOT_DIR = r'D:\tmp\A-tmp\models\classify-leaves'
TRAIN_CSV = os.path.join(ROOT_DIR, 'train.csv')
TEST_CSV = os.path.join(ROOT_DIR, 'test.csv')
BATCH_SIZE = 32
EPOCHS = 10
LR = 1e-4

# ==================== Step 1：数据 ====================
# -------------------- 数据预处理 --------------------
# 对训练数据做标签映射
train_df = pd.read_csv(TRAIN_CSV)
label_encoder = LabelEncoder()
# 直接对所有训练标签进行拟合
train_df['label_idx'] = label_encoder.fit_transform(train_df['label']) # 新增一列'label_id'，0-175
num_classes = len(label_encoder.classes_) # 176个类别
  
# 数据增强，transforms变换
train_transforms = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_transforms = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# -------------------- 自定义数据集类 --------------------
class CustomDataset(Dataset):
    def __init__(self, dataframe, root_dir, transform=None, is_test=False):
        self.data = dataframe
        self.root_dir = root_dir
        self.transform = transform
        self.is_test = is_test
  
    def __len__(self):
        return len(self.data)
  
    def __getitem__(self, idx):
        img_name = self.data.iloc[idx]['image']
        img_path = os.path.join(self.root_dir, img_name)
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
  
        if self.is_test:
            return image
        else:
            label = self.data.iloc[idx]['label_idx']
            label = torch.tensor(label, dtype=torch.long)
            return image, label

# -------------------- 实例化，用于后续训练和测试 --------------------
full_dataset = CustomDataset(train_df, ROOT_DIR, transform=train_transforms, is_test=False)
# 划分训练集和验证集 (8:2)
train_size = int(0.8 * len(full_dataset))
val_size = len(full_dataset) - train_size
train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])
# 测试集
test_df = pd.read_csv(TEST_CSV)
test_dataset = CustomDataset(test_df, ROOT_DIR, transform=test_transforms, is_test=True)
# DataLoader
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

# ==================== Step 2：模型 ====================
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, num_classes)
model = model.to(DEVICE)

# ==================== Step 3：损失函数 ====================
criterion = nn.CrossEntropyLoss()

# ==================== Step 4：优化器 ====================
optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=1e-4)
# 余弦退火学习率
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=1e-6)

# ==================== Step 5：迭代训练 ====================
def train_one_epoch(model, loader, criterion, optimizer):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    for images, labels in tqdm(loader, desc='Training'):
        images, labels = images.to(DEVICE), labels.to(DEVICE)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()
    return total_loss / len(loader), 100. * correct / total
  
def validate(model, loader, criterion):
    model.eval()
    total_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in tqdm(loader, desc='Validating'):
            images, labels = images.to(DEVICE), labels.to(DEVICE)
            outputs = model(images)
            loss = criterion(outputs, labels)
            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    return total_loss / len(loader), 100. * correct / total
  
# 训练循环
"""
引入 `best_acc` 动态保存模型，是工程实践中防止"训练过头"或"过拟合"的必备手段
"""
best_acc = 0
for epoch in range(EPOCHS):
    print(f'\nEpoch {epoch+1}/{EPOCHS}')
    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer)
    val_loss, val_acc = validate(model, val_loader, criterion)
    scheduler.step()
    print(f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%')
    print(f'Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%')
    # 保存最佳模型
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(model.state_dict(), 'best_model.pth')
        print(f'Model saved with acc: {val_acc:.2f}%')

# ==================== 测试预测 ====================
model.eval()
predictions = []
with torch.no_grad():
    for images in test_loader:
        images = images.to(DEVICE)
        outputs = model(images)
        _, predicted = outputs.max(1)
        predictions.extend(predicted.cpu().numpy())
  
# 使用 label_encoder 进行反向转换
predicted_labels = label_encoder.inverse_transform(predictions)
  
# 生成提交文件
submission = pd.read_csv(TEST_CSV)
submission['label'] = predicted_labels
submission.to_csv('submission_resnet18.csv', index=False)
print("恭喜！预测完成并已生成 submission.csv")

对"Step1：数据"这步，数据预处理相对来说是比较麻烦的，对于不同的任务相应的预处理也不尽相同，主要的几点是：1.对图片数据的标签用LabelEncoder做标签映射；2.自定义数据集CustomDataset；3.实例化数据集，划分好训练/验证/测试之后，利用trainloader进行迭代。

对"Step2：模型"这步用的是预训练好的ResNet18。

对"Step5：迭代训练"这步，有很多"套路"和行业标准写法。

标准的代码逻辑结构

train_one_epoch 和 validate 函数是一种模块化封装 。

好处：

①这样写逻辑清晰，训练模式和评估模式的切换非常明确；
②防止训练过程中的临时变量（如 outputs, loss）污染验证过程；
③方便复用

python 复制代码

for epoch in range(EPOCHS):
    # --- 训练阶段 ---
    train_loss, train_acc = train_one_epoch(...) # 开启 model.train(), optimizer.step()
    
    # --- 验证阶段 ---
    val_loss, val_acc = validate(...)           # 开启 model.eval(), torch.no_grad()
    
    # --- 策略调整 (套路的核心) ---
    scheduler.step()                             # 更新学习率
    
    # --- 模型持久化 (Checkpoint) ---
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(model.state_dict(), 'best_model.pth') # 只存最好的那个

tqdm

tqdm 只是一个进度条显示工具 。是对「可迭代对象」的 非侵入式 可视化封装。
作用：

1️⃣ 实时查看剩余时间（ETA）

2️⃣ 每秒迭代次数（it/s）

3️⃣ 当前的 Loss（或其他指标）

这些可视化的作用，可以方便我们判断程序有没有"卡死"；性能瓶颈；loss 是否"正常变化"

model.train()和model.eval()

在 PyTorch 中，model.train() 和 model.eval() 用于切换模型的运行模式（mode）

它们不会直接控制是否计算梯度 ，而是主要影响模型中具有不同行为的层（如 BatchNorm 和 Dropout）

model.train() ------ 训练模式

model.train() 用于开启训练模式，其作用包括：

将模型切换为 training mode
使模型中某些层（如 BatchNorm、Dropout）表现为训练时的行为
配合反向传播 （loss.backward()）完成参数更新

在 train() 模式下，BatchNorm 和 Dropout 行为都正常

model.eval() ------ 评估模式

model.eval() 用于开启评估 / 推理模式，其作用包括：

将模型切换为 evaluation mode
固定 BatchNorm、Dropout 等层的行为
通常与 torch.no_grad() 配合使用

在 eval() 模式下，BatchNorm不再使用当前 batch 的统计量，而是使用训练阶段累计得到的 running mean / running variance；Dropout 被完全关闭，所有神经元都会参与前向传播

总结：当模型中包含 BatchNorm / Dropout 层时，必须显式调用 model.train() 和 model.eval()。

⚠️ 当模型中不包含这些层时不调用可能"看起来能跑"，但这依赖于模型结构，不具备通用性和可扩展性。

总结

迭代训练套路模板

在 PyTorch 等主流深度学习框架中，模型迭代训练的核心套路模板本质上只有两种：

1️⃣ 内联写法：过程式编程，在 epoch 循环内直接写 train / eval 逻辑

2️⃣ 模块化写法：函数式/结构化编程，将 train / eval 封装成函数，在 epoch 循环中调用

是否存在"第三种本质套路"？

之后可能看到的：PyTorch Lightning、HuggingFace Trainer、Detectron2 / MMDetection......

本质上都是：把"模块化写法：函数式/结构化编程"进一步自动化和框架化 ，而不是新的训练逻辑。这些框架通过封装和抽象，将常见的训练流程和功能（如日志记录、模型保存、回调机制等）进行了标准化和自动化，从而简化了开发流程 。

比如所谓的Hooks / Callbacks（框架式训练）写法，通过定义一个Trainer类来实现。在这个类中，你可以定义训练和验证的逻辑，并通过回调机制在不同的训练阶段执行特定的操作（如保存模型、发送通知等）。