训练流程#

训练神经网络是深度学习的核心环节。本章将详细介绍 PyTorch 中完整的训练流程,包括前向传播、反向传播、优化器使用等。

基本训练循环#

最简单的训练循环#

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

# 1. 准备数据
# ... (假设已有 dataloader)

# 2. 定义模型
model = nn.Linear(10, 1)

# 3. 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 4. 训练循环
model.train()  # 设置为训练模式
for epoch in range(num_epochs):
    for batch_data, batch_labels in dataloader:
        # 前向传播
        outputs = model(batch_data)
        loss = criterion(outputs, batch_labels)
        
        # 反向传播
        optimizer.zero_grad()  # 清零梯度
        loss.backward()        # 计算梯度
        optimizer.step()       # 更新参数
        
        # 打印损失
        if (epoch + 1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

完整的训练和验证流程#

def train_model(model, train_loader, val_loader, num_epochs, device):
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    # 记录训练历史
    train_losses = []
    val_losses = []
    train_accs = []
    val_accs = []
    
    for epoch in range(num_epochs):
        # ========== 训练阶段 ==========
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        for batch_data, batch_labels in train_loader:
            # 移到设备
            batch_data = batch_data.to(device)
            batch_labels = batch_labels.to(device)
            
            # 前向传播
            outputs = model(batch_data)
            loss = criterion(outputs, batch_labels)
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            # 统计
            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_total += batch_labels.size(0)
            train_correct += (predicted == batch_labels).sum().item()
        
        # 计算平均损失和准确率
        avg_train_loss = train_loss / len(train_loader)
        train_acc = 100 * train_correct / train_total
        
        # ========== 验证阶段 ==========
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():  # 验证时不需要计算梯度
            for batch_data, batch_labels in val_loader:
                batch_data = batch_data.to(device)
                batch_labels = batch_labels.to(device)
                
                outputs = model(batch_data)
                loss = criterion(outputs, batch_labels)
                
                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                val_total += batch_labels.size(0)
                val_correct += (predicted == batch_labels).sum().item()
        
        avg_val_loss = val_loss / len(val_loader)
        val_acc = 100 * val_correct / val_total
        
        # 记录历史
        train_losses.append(avg_train_loss)
        val_losses.append(avg_val_loss)
        train_accs.append(train_acc)
        val_accs.append(val_acc)
        
        # 打印进度
        print(f'Epoch [{epoch+1}/{num_epochs}]')
        print(f'Train Loss: {avg_train_loss:.4f}, Train Acc: {train_acc:.2f}%')
        print(f'Val Loss: {avg_val_loss:.4f}, Val Acc: {val_acc:.2f}%')
        print('-' * 50)
    
    return {
        'train_losses': train_losses,
        'val_losses': val_losses,
        'train_accs': train_accs,
        'val_accs': val_accs
    }

优化器使用#

常用优化器#

# SGD
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)

# Adam
optimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999), weight_decay=1e-4)

# AdamW
optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=1e-4)

# RMSprop
optimizer = optim.RMSprop(model.parameters(), lr=0.01, alpha=0.99)

# Adagrad
optimizer = optim.Adagrad(model.parameters(), lr=0.01)

学习率调度器#

# StepLR:每 step_size 个 epoch 将学习率乘以 gamma
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

# ExponentialLR:每个 epoch 将学习率乘以 gamma
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95)

# CosineAnnealingLR:余弦退火
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

# ReduceLROnPlateau:当指标停止改善时降低学习率
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.5, patience=5
)

# 在训练循环中使用
for epoch in range(num_epochs):
    # ... 训练代码 ...
    
    # 更新学习率
    scheduler.step()  # 对于 StepLR, ExponentialLR, CosineAnnealingLR
    # 或
    scheduler.step(val_loss)  # 对于 ReduceLROnPlateau

不同参数组使用不同学习率#

optimizer = optim.SGD([
    {'params': model.features.parameters(), 'lr': 0.001},  # 预训练层
    {'params': model.classifier.parameters(), 'lr': 0.01}  # 新层
], lr=0.001, momentum=0.9)

梯度裁剪#

# 梯度裁剪:防止梯度爆炸
max_norm = 1.0
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

# 在训练循环中
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
optimizer.step()

早停(Early Stopping)#

class EarlyStopping:
    def __init__(self, patience=7, min_delta=0, restore_best_weights=True):
        self.patience = patience
        self.min_delta = min_delta
        self.restore_best_weights = restore_best_weights
        self.best_loss = None
        self.counter = 0
        self.best_weights = None
    
    def __call__(self, val_loss, model):
        if self.best_loss is None:
            self.best_loss = val_loss
            self.save_checkpoint(model)
        elif val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter = 0
            self.save_checkpoint(model)
        else:
            self.counter += 1
        
        if self.counter >= self.patience:
            if self.restore_best_weights:
                model.load_state_dict(self.best_weights)
            return True
        return False
    
    def save_checkpoint(self, model):
        self.best_weights = model.state_dict().copy()

# 使用
early_stopping = EarlyStopping(patience=10)
for epoch in range(num_epochs):
    # ... 训练和验证 ...
    if early_stopping(val_loss, model):
        print("Early stopping triggered")
        break

混合精度训练#

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for epoch in range(num_epochs):
    for batch_data, batch_labels in train_loader:
        optimizer.zero_grad()
        
        # 使用混合精度
        with autocast():
            outputs = model(batch_data)
            loss = criterion(outputs, batch_labels)
        
        # 缩放损失并反向传播
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

分布式训练#

DataParallel(单机多GPU)#

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

model = model.to(device)

DistributedDataParallel(多机多GPU,推荐)#

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化进程组
dist.init_process_group(backend='nccl')

# 创建模型
model = DDP(model, device_ids=[local_rank])

# 使用 DistributedSampler
from torch.utils.data.distributed import DistributedSampler
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

训练技巧#

1. 学习率预热(Warmup)#

from torch.optim.lr_scheduler import LambdaLR

def warmup_lambda(epoch):
    if epoch < 5:
        return epoch / 5
    return 1.0

scheduler = LambdaLR(optimizer, lr_lambda=warmup_lambda)

2. 梯度累积#

accumulation_steps = 4
optimizer.zero_grad()

for i, (data, labels) in enumerate(dataloader):
    outputs = model(data)
    loss = criterion(outputs, labels) / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

3. 模型检查点#

def save_checkpoint(model, optimizer, epoch, loss, filepath):
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss,
    }
    torch.save(checkpoint, filepath)

def load_checkpoint(model, optimizer, filepath):
    checkpoint = torch.load(filepath)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint['epoch']
    loss = checkpoint['loss']
    return epoch, loss

4. 训练监控#

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('runs/experiment_1')

for epoch in range(num_epochs):
    # ... 训练 ...
    writer.add_scalar('Loss/Train', train_loss, epoch)
    writer.add_scalar('Loss/Val', val_loss, epoch)
    writer.add_scalar('Accuracy/Train', train_acc, epoch)
    writer.add_scalar('Accuracy/Val', val_acc, epoch)
    writer.add_scalar('Learning_Rate', optimizer.param_groups[0]['lr'], epoch)

writer.close()

完整训练模板#

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

def train(model, train_loader, val_loader, config):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
    
    writer = SummaryWriter(config['log_dir'])
    best_val_acc = 0.0
    
    for epoch in range(config['num_epochs']):
        # 训练
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        for data, labels in train_loader:
            data, labels = data.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(data)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_total += labels.size(0)
            train_correct += (predicted == labels).sum().item()
        
        # 验证
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            for data, labels in val_loader:
                data, labels = data.to(device), labels.to(device)
                outputs = model(data)
                loss = criterion(outputs, labels)
                
                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()
        
        # 记录
        train_acc = 100 * train_correct / train_total
        val_acc = 100 * val_correct / val_total
        
        writer.add_scalar('Loss/Train', train_loss / len(train_loader), epoch)
        writer.add_scalar('Loss/Val', val_loss / len(val_loader), epoch)
        writer.add_scalar('Accuracy/Train', train_acc, epoch)
        writer.add_scalar('Accuracy/Val', val_acc, epoch)
        
        scheduler.step()
        
        # 保存最佳模型
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), config['save_path'])
        
        print(f'Epoch [{epoch+1}/{config["num_epochs"]}]')
        print(f'Train Acc: {train_acc:.2f}%, Val Acc: {val_acc:.2f}%')
    
    writer.close()
    return model

常见问题#

1. 过拟合#

# 解决方案:
# - 增加 Dropout
# - 使用数据增强
# - 增加权重衰减(weight_decay)
# - 使用早停

2. 训练不收敛#

# 解决方案:
# - 检查学习率(可能太大或太小)
# - 检查数据预处理
# - 检查损失函数
# - 使用梯度裁剪

3. 内存不足#

# 解决方案:
# - 减小 batch_size
# - 使用梯度累积
# - 使用混合精度训练
# - 减少模型大小

练习#

  1. 实现一个完整的训练循环,包括训练和验证
  2. 使用学习率调度器优化训练过程
  3. 实现早停机制
  4. 使用 TensorBoard 可视化训练过程