Chuoer47's Blog

Appearance

模型训练与评估 ​

PyTorch 提供了灵活而强大的训练框架。本文将系统地介绍从训练循环搭建、评估指标实现、过拟合诊断到 TensorBoard 可视化的完整流程，并通过一个 CIFAR-10 图像分类案例串联所有知识点。

完整训练循环模板 ​

PyTorch 的训练循环是所有深度学习任务的核心骨架，包含以下五个关键步骤：forward（前向传播）→ loss（计算损失）→ backward（反向传播）→ step（更新参数）→ zero_grad（清零梯度）。

基本训练循环 ​

python
import torch
import torch.nn as nn
import torch.optim as optim

# 1. 定义模型、损失函数、优化器
model = MyModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 2. 训练循环
num_epochs = 10
for epoch in range(num_epochs):
    model.train()  # 设置为训练模式
    running_loss = 0.0

    for batch_idx, (inputs, targets) in enumerate(train_loader):
        # ---- forward ----
        outputs = model(inputs)

        # ---- loss ----
        loss = criterion(outputs, targets)

        # ---- backward ----
        optimizer.zero_grad()  # 清零梯度（重要！放在 backward 之前或之后均可）
        loss.backward()        # 反向传播，计算梯度

        # ---- step ----
        optimizer.step()       # 更新模型参数

        running_loss += loss.item()

    avg_loss = running_loss / len(train_loader)
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}")

zero_grad 放在哪里？ ​

optimizer.zero_grad() 的位置有两种常见写法，功能等价：

python
# 写法一：在 backward 之前清零（推荐，语义更清晰）
optimizer.zero_grad()
loss.backward()
optimizer.step()

# 写法二：在 step 之后清零
loss.backward()
optimizer.step()
optimizer.zero_grad()

注意：如果不清零梯度，梯度会在多次 backward 中累加。这在某些场景（如梯度累积）中是有意为之的，但在常规训练中必须每次清零。

带验证的完整训练循环 ​

python
for epoch in range(num_epochs):
    # ---------- 训练阶段 ----------
    model.train()
    train_loss = 0.0
    train_correct = 0
    train_total = 0

    for inputs, targets in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        train_total += targets.size(0)
        train_correct += predicted.eq(targets).sum().item()

    train_acc = 100.0 * train_correct / train_total

    # ---------- 验证阶段 ----------
    model.eval()
    val_loss = 0.0
    val_correct = 0
    val_total = 0

    with torch.no_grad():
        for inputs, targets in val_loader:
            outputs = model(inputs)
            loss = criterion(outputs, targets)

            val_loss += loss.item()
            _, predicted = outputs.max(1)
            val_total += targets.size(0)
            val_correct += predicted.eq(targets).sum().item()

    val_acc = 100.0 * val_correct / val_total

    print(f"Epoch [{epoch+1}/{num_epochs}] "
          f"Train Loss: {train_loss/len(train_loader):.4f} "
          f"Train Acc: {train_acc:.2f}% "
          f"Val Loss: {val_loss/len(val_loader):.4f} "
          f"Val Acc: {val_acc:.2f}%")

训练模式与评估模式 ​

PyTorch 中部分模块在训练和评估时行为不同，通过 model.train() 和 model.eval() 切换。

两种模式的差异 ​

特性model.train()model.eval()
Dropout随机丢弃神经元所有神经元参与计算
BatchNorm使用当前 batch 的均值/方差使用训练时累积的运行均值/方差
梯度计算计算梯度需配合 torch.no_grad()

受影响的常见模块 ​

python
# Dropout —— 训练时随机置零，评估时全部保留
drop = nn.Dropout(p=0.5)
drop.train()
print(drop(torch.ones(10)))   # 大约一半被置零

drop.eval()
print(drop(torch.ones(10)))   # 全部保留（但值变为 1-p 的缩放后值，eval 下不做缩放）

# BatchNorm —— 训练时用 batch 统计量，评估时用全局统计量
bn = nn.BatchNorm1d(100)
bn.train()   # 使用当前 batch 的 mean / var
bn.eval()     # 使用 running_mean / running_var

常见陷阱 ​

python
# 错误：忘记切换模式导致验证结果不稳定
model.eval()  # 忘记调用，Dropout 和 BatchNorm 仍在训练模式
with torch.no_grad():
    outputs = model(val_inputs)

# 正确做法
model.eval()
with torch.no_grad():
    outputs = model(val_inputs)

最佳实践：在验证和测试代码之前，始终显式调用 model.eval()，结束后用 model.train() 切回。

验证与评估流程 ​

完整评估函数 ​

python
def evaluate(model, dataloader, criterion, device='cpu'):
    """通用评估函数"""
    model.eval()
    total_loss = 0.0
    all_preds = []
    all_targets = []
    all_probs = []

    with torch.no_grad():
        for inputs, targets in dataloader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)

            total_loss += loss.item() * inputs.size(0)
            probs = torch.softmax(outputs, dim=1)
            _, predicted = probs.max(1)

            all_preds.extend(predicted.cpu().numpy())
            all_targets.extend(targets.cpu().numpy())
            all_probs.extend(probs.cpu().numpy())

    avg_loss = total_loss / len(dataloader.dataset)
    return avg_loss, all_preds, all_targets, all_probs

使用模型进行推理 ​

python
def predict(model, image, device='cpu'):
    """单张图片推理"""
    model.eval()
    with torch.no_grad():
        image = image.unsqueeze(0).to(device)  # 添加 batch 维度
        output = model(image)
        prob = torch.softmax(output, dim=1)
        predicted_class = prob.argmax(dim=1).item()
        confidence = prob.max().item()
    return predicted_class, confidence

常用评估指标实现 ​

分类指标 ​

Accuracy（准确率） ​

python
def accuracy(outputs, targets):
    """计算准确率"""
    _, predicted = outputs.max(1)
    correct = predicted.eq(targets).sum().item()
    return correct / targets.size(0)

Precision、Recall、F1-Score ​

python
def precision_recall_f1(outputs, targets, num_classes, average='macro'):
    """
    计算 Precision、Recall、F1-Score
    average: 'macro'（各类别等权平均）, 'micro'（全局计算）, 'weighted'（按类别频率加权）
    """
    _, predicted = outputs.max(1)

    precision_list = []
    recall_list = []
    f1_list = []

    for cls in range(num_classes):
        tp = ((predicted == cls) & (targets == cls)).sum().float()
        fp = ((predicted == cls) & (targets != cls)).sum().float()
        fn = ((predicted != cls) & (targets == cls)).sum().float()

        precision = tp / (tp + fp + 1e-8)
        recall = tp / (tp + fn + 1e-8)
        f1 = 2 * precision * recall / (precision + recall + 1e-8)

        precision_list.append(precision.item())
        recall_list.append(recall.item())
        f1_list.append(f1.item())

    if average == 'macro':
        return sum(precision_list) / num_classes, sum(recall_list) / num_classes, sum(f1_list) / num_classes
    elif average == 'micro':
        # 微平均：将所有类别的 TP/FP/FN 汇总后计算
        pass  # 逻辑类似，汇总后一次性计算

Confusion Matrix（混淆矩阵） ​

python
import numpy as np

def confusion_matrix(outputs, targets, num_classes):
    """计算混淆矩阵"""
    _, predicted = outputs.max(1)
    cm = np.zeros((num_classes, num_classes), dtype=np.int64)
    for t, p in zip(targets.numpy(), predicted.numpy()):
        cm[t][p] += 1
    return cm

def plot_confusion_matrix(cm, class_names):
    """可视化混淆矩阵"""
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(8, 8))
    im = ax.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
    ax.set_title('Confusion Matrix')
    plt.colorbar(im, ax=ax)
    tick_marks = np.arange(len(class_names))
    ax.set_xticks(tick_marks)
    ax.set_xticklabels(class_names, rotation=45)
    ax.set_yticks(tick_marks)
    ax.set_yticklabels(class_names)
    ax.set_ylabel('True Label')
    ax.set_xlabel('Predicted Label')
    # 在格子中显示数值
    for i in range(len(class_names)):
        for j in range(len(class_names)):
            ax.text(j, i, str(cm[i][j]), ha='center', va='center',
                    color='white' if cm[i][j] > cm.max() / 2 else 'black')
    plt.tight_layout()
    plt.show()

AUC-ROC ​

python
from sklearn.metrics import roc_auc_score, roc_curve
import matplotlib.pyplot as plt

def compute_auc_roc(all_probs, all_targets, num_classes):
    """计算多分类 AUC-ROC（One-vs-Rest）"""
    from sklearn.preprocessing import label_binarize
    targets_onehot = label_binarize(all_targets, classes=list(range(num_classes)))
    probs = np.array(all_probs)

    # 宏平均 AUC
    auc = roc_auc_score(targets_onehot, probs, multi_class='ovr', average='macro')
    return auc

def plot_roc_curve(all_probs, all_targets, num_classes, class_names):
    """绘制 ROC 曲线"""
    from sklearn.preprocessing import label_binarize
    targets_onehot = label_binarize(all_targets, classes=list(range(num_classes)))
    probs = np.array(all_probs)

    plt.figure(figsize=(10, 8))
    for i in range(num_classes):
        fpr, tpr, _ = roc_curve(targets_onehot[:, i], probs[:, i])
        auc_val = roc_auc_score(targets_onehot[:, i], probs[:, i])
        plt.plot(fpr, tpr, label=f'{class_names[i]} (AUC={auc_val:.3f})')

    plt.plot([0, 1], [0, 1], 'k--', label='Random')
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('ROC Curve (One-vs-Rest)')
    plt.legend()
    plt.grid(True)
    plt.show()

回归指标 ​

python
import torch.nn.functional as F

def mse(outputs, targets):
    """均方误差"""
    return F.mse_loss(outputs, targets).item()

def rmse(outputs, targets):
    """均方根误差"""
    return torch.sqrt(F.mse_loss(outputs, targets)).item()

def mae(outputs, targets):
    """平均绝对误差"""
    return F.l1_loss(outputs, targets).item()

def r_squared(outputs, targets):
    """决定系数 R²"""
    ss_res = ((targets - outputs) ** 2).sum()
    ss_tot = ((targets - targets.mean()) ** 2).sum()
    return (1 - ss_res / (ss_tot + 1e-8)).item()

# 使用示例
y_pred = torch.tensor([2.5, 0.0, 2.1, 7.8])
y_true = torch.tensor([3.0, -0.5, 2.0, 7.5])
print(f"MSE:  {mse(y_pred, y_true):.4f}")
print(f"RMSE: {rmse(y_pred, y_true):.4f}")
print(f"MAE:  {mae(y_pred, y_true):.4f}")
print(f"R²:   {r_squared(y_pred, y_true):.4f}")

过拟合与欠拟合诊断 ​

学习曲线分析 ​

通过观察训练集和验证集的 loss/accuracy 变化，可以判断模型状态：

现象诊断解决策略
训练 loss 高，验证 loss 高欠拟合增大模型容量、降低正则化、延长训练
训练 loss 低，验证 loss 高且差距大过拟合增加正则化、数据增强、减小模型容量
训练 loss 低，验证 loss 低且差距小良好拟合模型状态理想

学习曲线绘制 ​

python
import matplotlib.pyplot as plt

def plot_learning_curves(train_losses, val_losses, train_accs=None, val_accs=None):
    """绘制学习曲线"""
    epochs = range(1, len(train_losses) + 1)

    fig, axes = plt.subplots(1, 2 if train_accs else 1, figsize=(14, 5))
    if not isinstance(axes, (list, np.ndarray)):
        axes = [axes]

    # Loss 曲线
    axes[0].plot(epochs, train_losses, 'b-', label='Train Loss')
    axes[0].plot(epochs, val_losses, 'r-', label='Val Loss')
    axes[0].set_xlabel('Epoch')
    axes[0].set_ylabel('Loss')
    axes[0].set_title('Loss Curve')
    axes[0].legend()
    axes[0].grid(True)

    # Accuracy 曲线（如果有）
    if train_accs and val_accs:
        axes[1].plot(epochs, train_accs, 'b-', label='Train Acc')
        axes[1].plot(epochs, val_accs, 'r-', label='Val Acc')
        axes[1].set_xlabel('Epoch')
        axes[1].set_ylabel('Accuracy (%)')
        axes[1].set_title('Accuracy Curve')
        axes[1].legend()
        axes[1].grid(True)

    plt.tight_layout()
    plt.show()

正则化策略 ​

Dropout ​

python
class RegularizedModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, dropout_rate=0.5):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(p=dropout_rate),      # 全连接层后添加 Dropout
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(p=dropout_rate),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        return self.net(x)

# 不同层可以设置不同 dropout 率
# 卷积网络中常用 nn.Dropout2d，按通道置零
conv_dropout = nn.Dropout2d(p=0.25)

Weight Decay（L2 正则化） ​

python
# 在优化器中设置 weight_decay 参数
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)

# 注意：通常不对 bias 和 BatchNorm 的参数施加 weight decay
def get_param_groups(model, weight_decay):
    """区分需要和不需要 weight decay 的参数"""
    decay_params = []
    no_decay_params = []

    for name, param in model.named_parameters():
        if 'bias' in name or 'bn' in name or 'norm' in name:
            no_decay_params.append(param)
        else:
            decay_params.append(param)

    return [
        {'params': decay_params, 'weight_decay': weight_decay},
        {'params': no_decay_params, 'weight_decay': 0.0}
    ]

param_groups = get_param_groups(model, weight_decay=1e-4)
optimizer = optim.Adam(param_groups, lr=1e-3)

数据增强 ​

python
from torchvision import transforms

# 训练集使用数据增强
train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.RandomRotation(degrees=15),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.4914, 0.4822, 0.4465],
                         std=[0.2470, 0.2435, 0.2616]),
    transforms.RandomErasing(p=0.25),  # CutOut 变体，在 ToTensor 之后
])

# 验证集不做增强
val_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.4914, 0.4822, 0.4465],
                         std=[0.2470, 0.2435, 0.2616]),
])

Early Stopping ​

python
class EarlyStopping:
    """早停机制：验证指标不再改善时提前终止训练"""

    def __init__(self, patience=10, min_delta=0.0, mode='min', verbose=True):
        """
        Args:
            patience: 容忍多少个 epoch 没有改善
            min_delta: 最小改善幅度
            mode: 'min'（指标越小越好，如 loss）或 'max'（越大越好，如 accuracy）
            verbose: 是否打印提示信息
        """
        self.patience = patience
        self.min_delta = min_delta
        self.mode = mode
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.best_model_state = None

    def __call__(self, score, model):
        if self.best_score is None:
            self.best_score = score
            self.best_model_state = model.state_dict().copy()
        elif self._is_improvement(score):
            self.best_score = score
            self.best_model_state = model.state_dict().copy()
            self.counter = 0
            if self.verbose:
                print(f"  Validation improved. Saving model...")
        else:
            self.counter += 1
            if self.verbose:
                print(f"  No improvement for {self.counter}/{self.patience} epochs.")
            if self.counter >= self.patience:
                self.early_stop = True
                if self.verbose:
                    print("  Early stopping triggered!")

    def _is_improvement(self, score):
        if self.mode == 'min':
            return score < self.best_score - self.min_delta
        else:
            return score > self.best_score + self.min_delta

    def load_best_model(self, model):
        """恢复最优模型权重"""
        model.load_state_dict(self.best_model_state)
        return model

# 使用示例
early_stopping = EarlyStopping(patience=10, mode='max', verbose=True)

for epoch in range(num_epochs):
    train_one_epoch(model, train_loader, optimizer, criterion)
    val_acc = evaluate(model, val_loader, criterion)

    early_stopping(val_acc, model)
    if early_stopping.early_stop:
        print(f"Early stopping at epoch {epoch+1}")
        break

# 恢复最优模型
model = early_stopping.load_best_model(model)

TensorBoard 可视化 ​

TensorBoard 是训练过程可视化的强大工具，PyTorch 通过 torch.utils.tensorboard 提供原生支持。

初始化 SummaryWriter ​

python
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(log_dir='runs/experiment_1')

记录 Loss 和 Accuracy 曲线 ​

python
for epoch in range(num_epochs):
    train_loss = train_one_epoch(...)
    val_loss, val_acc = evaluate(...)

    # 使用 add_scalar 记录标量
    writer.add_scalar('Loss/train', train_loss, epoch)
    writer.add_scalar('Loss/val', val_loss, epoch)
    writer.add_scalar('Accuracy/val', val_acc, epoch)

    # 学习率记录
    writer.add_scalar('LR', optimizer.param_groups[0]['lr'], epoch)

writer.close()

在同一图表中对比多个实验 ​

python
# 使用不同的 log_dir 或 add_hparams
for lr in [1e-2, 1e-3, 1e-4]:
    writer = SummaryWriter(f'runs/lr_{lr}')
    # ... 训练代码 ...
    writer.close()

记录模型图 ​

python
# 将模型结构写入 TensorBoard
dummy_input = torch.randn(1, 3, 32, 32)
writer.add_graph(model, dummy_input)
writer.close()

记录图像 ​

python
from torchvision.utils import make_grid

# 记录一个 batch 的训练图片
dataiter = iter(train_loader)
images, labels = next(dataiter)

# add_images 会自动将多张图片拼成网格
writer.add_image('Train Images', make_grid(images, nrow=8), global_step=0)

# 记录模型的卷积核权重
def log_conv_weights(model, writer, epoch):
    for name, param in model.named_parameters():
        if 'conv' in name and 'weight' in name:
            # 取前 16 个卷积核
            kernels = param[:16].detach().cpu()
            # 归一化到 [0, 1]
            kernels = (kernels - kernels.min()) / (kernels.max() - kernels.min() + 1e-8)
            writer.add_images(f'Weights/{name}', kernels.unsqueeze(1), epoch)

记录直方图 ​

python
def log_histograms(model, writer, epoch):
    """记录参数和梯度的分布"""
    for name, param in model.named_parameters():
        if param.grad is not None:
            writer.add_histogram(f'Gradients/{name}', param.grad, epoch)
        writer.add_histogram(f'Parameters/{name}', param.data, epoch)

启动 TensorBoard ​

bash
# 在终端运行
tensorboard --logdir=runs --port=6006

# 然后在浏览器中访问 http://localhost:6006

完整图像分类训练案例（CIFAR-10） ​

以下是一个完整的、可直接运行的 CIFAR-10 图像分类训练代码：

python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
import time

# ==================== 1. 超参数配置 ====================
class Config:
    batch_size = 128
    num_epochs = 50
    learning_rate = 0.001
    weight_decay = 1e-4
    num_classes = 10
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    log_dir = 'runs/cifar10_resnet'

config = Config()

# ==================== 2. 数据准备 ====================
train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
])

test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
])

train_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=train_transform)
test_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=test_transform)

train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=config.batch_size, shuffle=True, num_workers=4)
test_loader = torch.utils.data.DataLoader(
    test_dataset, batch_size=config.batch_size, shuffle=False, num_workers=4)

classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

# ==================== 3. 模型定义（简化版 ResNet） ====================
class BasicBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        return self.relu(out)

class SimpleResNet(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, 1, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)

        self.layer1 = self._make_layer(64, 64, 2, stride=1)
        self.layer2 = self._make_layer(64, 128, 2, stride=2)
        self.layer3 = self._make_layer(128, 256, 2, stride=2)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(256, num_classes)

    def _make_layer(self, in_ch, out_ch, num_blocks, stride):
        layers = [BasicBlock(in_ch, out_ch, stride)]
        for _ in range(1, num_blocks):
            layers.append(BasicBlock(out_ch, out_ch))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return self.fc(x)

# ==================== 4. 训练与评估 ====================
model = SimpleResNet(config.num_classes).to(config.device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=config.learning_rate,
                       weight_decay=config.weight_decay)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.num_epochs)
writer = SummaryWriter(config.log_dir)

best_acc = 0.0

for epoch in range(config.num_epochs):
    # ---- 训练 ----
    model.train()
    train_loss = 0.0
    train_correct = 0
    train_total = 0
    start_time = time.time()

    for images, labels in train_loader:
        images, labels = images.to(config.device), labels.to(config.device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        train_total += labels.size(0)
        train_correct += predicted.eq(labels).sum().item()

    scheduler.step()
    train_acc = 100.0 * train_correct / train_total
    avg_train_loss = train_loss / len(train_loader)

    # ---- 测试 ----
    model.eval()
    test_loss = 0.0
    test_correct = 0
    test_total = 0

    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(config.device), labels.to(config.device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            test_loss += loss.item()
            _, predicted = outputs.max(1)
            test_total += labels.size(0)
            test_correct += predicted.eq(labels).sum().item()

    test_acc = 100.0 * test_correct / test_total
    avg_test_loss = test_loss / len(test_loader)
    elapsed = time.time() - start_time

    # ---- 日志 ----
    writer.add_scalar('Loss/train', avg_train_loss, epoch)
    writer.add_scalar('Loss/test', avg_test_loss, epoch)
    writer.add_scalar('Accuracy/train', train_acc, epoch)
    writer.add_scalar('Accuracy/test', test_acc, epoch)
    writer.add_scalar('LR', scheduler.get_last_lr()[0], epoch)

    print(f"Epoch [{epoch+1}/{config.num_epochs}] "
          f"Train Loss: {avg_train_loss:.4f} Acc: {train_acc:.2f}% | "
          f"Test Loss: {avg_test_loss:.4f} Acc: {test_acc:.2f}% | "
          f"Time: {elapsed:.1f}s")

    # ---- 保存最优模型 ----
    if test_acc > best_acc:
        best_acc = test_acc
        torch.save(model.state_dict(), 'best_model.pth')
        print(f"  => Saved best model with accuracy: {best_acc:.2f}%")

writer.close()
print(f"Training complete. Best accuracy: {best_acc:.2f}%")

训练技巧汇总 ​

Learning Rate Warmup ​

python
class WarmupCosineScheduler:
    """带 Warmup 的余弦退火调度器"""

    def __init__(self, optimizer, warmup_epochs, total_epochs, base_lr, min_lr=1e-6):
        self.optimizer = optimizer
        self.warmup_epochs = warmup_epochs
        self.total_epochs = total_epochs
        self.base_lr = base_lr
        self.min_lr = min_lr

    def step(self, epoch):
        import math
        if epoch < self.warmup_epochs:
            # 线性 warmup
            lr = self.base_lr * (epoch + 1) / self.warmup_epochs
        else:
            # 余弦退火
            progress = (epoch - self.warmup_epochs) / (self.total_epochs - self.warmup_epochs)
            lr = self.min_lr + 0.5 * (self.base_lr - self.min_lr) * (1 + math.cos(math.pi * progress))

        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
        return lr

Gradient Accumulation ​

当显存不足以容纳大 batch 时，通过梯度累积模拟大 batch 训练：

python
accumulation_steps = 4  # 累积 4 个 mini-batch 的梯度，等效于 4 倍 batch_size

optimizer.zero_grad()
for i, (inputs, targets) in enumerate(train_loader):
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss = loss / accumulation_steps  # 归一化 loss
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# 处理最后不足 accumulation_steps 的部分
if (i + 1) % accumulation_steps != 0:
    optimizer.step()
    optimizer.zero_grad()

Mixed Precision Training（混合精度训练） ​

混合精度训练使用 FP16 加速计算，同时保持 FP32 的关键操作以确保数值稳定性：

python
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
model = model.cuda()

for epoch in range(num_epochs):
    model.train()
    for inputs, targets in train_loader:
        inputs, targets = inputs.cuda(), targets.cuda()

        optimizer.zero_grad()

        # 自动混合精度前向传播
        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, targets)

        # 缩放 loss 后反向传播
        scaler.scale(loss).backward()

        # 先 unscale 再裁剪梯度（可选）
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        # 更新参数
        scaler.step(optimizer)
        scaler.update()  # 更新缩放因子

PyTorch 2.x torch.compile 加速 ​

python
# PyTorch 2.0+ 的编译优化，可显著加速训练
model = torch.compile(model, mode='reduce-overhead')
# mode 选项：
#   'default'      - 默认优化
#   'reduce-overhead' - 减少 Python 开销，适合小 batch
#   'max-autotune' - 最大化优化，编译较慢但运行最快

参考资料 ​