从零到一跑通第一个训练脚本

完整的训练循环、损失函数/优化器选择、训练模式切换、梯度累积等核心内容。

[toc]

训练流程全解析

---

🔄 完整训练循环

for epoch in range(num_epochs):
    # ===== 训练阶段 =====
    model.train()
    for batch_x, batch_y in train_loader:
        # 1. 前向传播
        outputs = model(batch_x)
        
        # 2. 计算损失
        loss = criterion(outputs, batch_y)
        
        # 3. 反向传播
        optimizer.zero_grad()  # 清除旧梯度
        loss.backward()        # 计算梯度
        
        # 4. 更新参数
        optimizer.step()       # 参数更新
    
    # ===== 验证阶段 =====
    model.eval()
    with torch.no_grad():
        for batch_x, batch_y in val_loader:
            outputs = model(batch_x)
            val_loss = criterion(outputs, batch_y)

---

1️⃣ 完整训练脚本 (模板)

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
import numpy as np
 
def train_one_epoch(model, dataloader, criterion, optimizer, device):
    """训练一个 epoch"""
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    
    pbar = tqdm(dataloader, desc="Training")
    for batch_x, batch_y in pbar:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        
        # 前向传播
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # 统计
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += batch_y.size(0)
        correct += predicted.eq(batch_y).sum().item()
        
        pbar.set_postfix({
            'loss': f'{loss.item():.4f}',
            'acc': f'{100.*correct/total:.2f}%'
        })
    
    return total_loss / len(dataloader), 100. * correct / total
 
 
def validate(model, dataloader, criterion, device):
    """验证一个 epoch"""
    model.eval()
    total_loss = 0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for batch_x, batch_y in tqdm(dataloader, desc="Validating"):
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            
            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += batch_y.size(0)
            correct += predicted.eq(batch_y).sum().item()
    
    return total_loss / len(dataloader), 100. * correct / total
 
 
def main():
    # 超参数
    BATCH_SIZE = 64
    EPOCHS = 50
    LEARNING_RATE = 1e-3
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 准备数据（以 MNIST 为例）
    from torchvision import datasets, transforms
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    train_dataset = datasets.MNIST(root='./data', train=True, 
                                   download=True, transform=transform)
    val_dataset = datasets.MNIST(root='./data', train=False, 
                                  download=True, transform=transform)
    
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, 
                              shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, 
                            shuffle=False, num_workers=4)
    
    # 创建模型
    from model import SimpleCNN
    model = SimpleCNN(num_classes=10).to(DEVICE)
    
    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=3
    )
    
    # 训练循环
    best_acc = 0
    for epoch in range(1, EPOCHS + 1):
        print(f"\nEpoch {epoch}/{EPOCHS}")
        
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, DEVICE
        )
        val_loss, val_acc = validate(
            model, val_loader, criterion, DEVICE
        )
        
        scheduler.step(val_loss)
        
        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"✅ 新最佳模型保存! Acc: {best_acc:.2f}%")
 
if __name__ == "__main__":
    main()

---

2️⃣ 核心组件的选择策略

损失函数

# 分类任务
criterion = nn.CrossEntropyLoss()           # 多分类 (最常用)
criterion = nn.BCEWithLogitsLoss()          # 多标签分类
criterion = nn.BCELoss()                    # 二分类 (配合 sigmoid)
 
# 回归任务
criterion = nn.MSELoss()                    # 均方误差
criterion = nn.L1Loss()                     # 平均绝对误差
 
# 特殊任务
criterion = nn.KLDivLoss()                  # 蒸馏/分布匹配
criterion = nn.CTCLoss()                    # 语音/OCR

优化器

# 新手推荐
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
 
# 进阶
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = torch.optim.RMSprop(model.parameters(), lr=1e-3)

学习率调度器

# 常用调度器
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.5, patience=5
)
 
# 预热 + 余弦退火 (训练 LLM 常见)
scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer, 
    lr_lambda=lambda epoch: min(epoch / warmup_epochs, 1.0) 
                           if epoch < warmup_epochs 
                           else 0.5 * (1 + cos(pi * epoch / total_epochs))
)

---

3️⃣ 训练模式 vs 评估模式

model.train()
# BatchNorm 使用当前 batch 统计量
# Dropout 生效，随机丢弃神经元
 
model.eval()
# BatchNorm 使用训练集积累的统计量
# Dropout **关闭**，全部神经元参与

⚠️ 忘记切换 eval 模式是新手最常见的 bug！

---

4️⃣ 梯度累积 (大模型必备技巧)

当显存不够，又想要大 batch size 时：

accumulation_steps = 4  # 等效 batch size = 32 * 4 = 128
optimizer.zero_grad()
 
for i, (inputs, labels) in enumerate(train_loader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps  # 归一化
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

---

🔗 下一篇

→ 超参数调优