训练流程#
训练神经网络是深度学习的核心环节。本章将详细介绍 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, loss4. 训练监控#
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
# - 使用梯度累积
# - 使用混合精度训练
# - 减少模型大小练习#
- 实现一个完整的训练循环,包括训练和验证
- 使用学习率调度器优化训练过程
- 实现早停机制
- 使用 TensorBoard 可视化训练过程