学习笔记 深度学习

PyTorch训练技巧与高级功能

加速训练的实用技术 -- 从混合精度到分布式训练的完整指南

一、概述

在实际的深度学习项目开发中，模型架构设计只是工作的一部分。如何高效、稳定地训练模型，往往决定了项目的成败。PyTorch 提供了丰富的训练工具和高级功能，帮助开发者应对梯度爆炸、显存不足、收敛缓慢、训练耗时过长等常见问题。本文系统梳理了八项核心训练优化技术，涵盖混合精度训练、梯度累积、梯度裁剪、学习率调度、EMA指数移动平均、分布式训练、实验记录与模型统计，每项技术均配有可直接运行的代码示例与实践建议。

核心内容一览

  • 混合精度训练： 利用FP16半精度计算加速训练、降低显存占用，同时保持模型精度
  • 梯度累积： 通过累积多个小批量的梯度来模拟大批量训练，突破显存限制
  • 梯度裁剪： 限制梯度范数或值域，有效防止梯度爆炸导致训练崩溃
  • 学习率调度： 动态调整学习率以优化收敛过程，避免陷入局部最优
  • EMA指数移动平均： 对模型参数做平滑处理，提升测试时模型泛化能力
  • 分布式训练： 利用多GPU并行计算，大幅缩短训练时间
  • 实验记录： 使用TensorBoard可视化训练指标，辅助调参决策
  • 模型统计： 使用torchinfo快速了解模型参数量、结构等信息

二、混合精度训练 (Mixed Precision Training)

2.1 基本原理

混合精度训练的核心思路是以FP16（半精度浮点数）执行前向传播和梯度计算，同时以FP32（单精度浮点数）维护权重主副本。FP16相比FP32将显存占用减半，并且在支持Tensor Core的GPU（如NVIDIA Volta及之后架构）上可以大幅提升矩阵运算吞吐量。然而FP16的数值范围较窄（约5.96e-8 ~ 65504），容易出现下溢（underflow）或上溢（overflow）问题。PyTorch的 torch.cuda.amp 模块通过 autocast 和 GradScaler 两个组件解决这一问题：autocast 自动为不同算子选择合适的精度，GradScaler 通过动态缩放损失值来防止梯度下溢。

为什么需要混合精度而非纯FP16？

纯FP16训练存在两个严重缺陷：一是梯度值小于FP16最小正数（约6e-8）时会直接变为0，导致参数无法更新；二是某些操作（如逐位求和、softmax等）在FP16下精度损失过大，影响模型收敛。混合精度策略保留FP32路径处理这些敏感操作，仅在安全的计算路径上使用FP16，实现了速度与精度的最佳平衡。

2.2 标准实现方式

# 混合精度训练完整示例 import torch import torch.nn as nn from torch.cuda.amp import autocast, GradScaler model = MyModel().cuda() optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) # 初始化GradScaler用于动态损失缩放 scaler = GradScaler() for batch in dataloader: inputs, labels = batch inputs, labels = inputs.cuda(), labels.cuda() optimizer.zero_grad() # 在autocast上下文内执行前向传播和损失计算 with autocast(): outputs = model(inputs) loss = nn.CrossEntropyLoss()(outputs, labels) # scaler.scale()替代loss.backward() scaler.scale(loss).backward() # scaler.step()替代optimizer.step() scaler.step(optimizer) # 更新缩放因子，准备下一次迭代 scaler.update()

2.3 关键参数与实践建议

GradScaler 在训练开始时使用较大的缩放因子（默认2^16 = 65536），每次 scaler.step(optimizer) 后检查是否有梯度溢出。若发现梯度为inf或NaN，则跳过本次参数更新并将缩放因子减半；若连续多次（growth_interval，默认200步）未出现溢出，则将缩放因子加倍。这一动态调节机制确保了训练的稳定性。

实践建议

  • 混合精度训练在batch size较大时加速效果更明显，通常可获得1.5x-3x的加速
  • 推荐与梯度累积配合使用，进一步降低单步显存峰值
  • NVIDIA Ampere（A100/A30）及以上架构支持BF16，其动态范围与FP32相当，可减少下溢风险
  • 使用 torch.set_autocast_enabled(True) 或环境变量 TORCH_AUTOCAST_ENABLED=1 全局启用

2.4 多卡场景下的混合精度

# DDP + 混合精度训练 from torch.nn.parallel import DistributedDataParallel as DDP from torch.cuda.amp import autocast, GradScaler model = DDP(model, device_ids=[local_rank]) scaler = GradScaler() for batch in dataloader: optimizer.zero_grad() with autocast(): outputs = model(inputs) loss = loss_fn(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

三、梯度累积 (Gradient Accumulation)

3.1 为什么需要梯度累积

在实际训练中，GPU显存往往限制了batch size的上限。然而较大的batch size对梯度估计更准确、Batch Normalization统计量更稳定、训练收敛更平滑。梯度累积的核心思想是：将一个"虚拟大批次"拆分为多个"实际小批次"逐个计算梯度，但不立即更新参数，而是累积这些梯度，待累积步数达到设定值后再统一执行 optimizer.step() 和 optimizer.zero_grad()。这等效于使用了 virtual_batch_size = physical_batch_size x accumulation_steps 进行训练。

3.2 标准实现

# 梯度累积实现 model = MyModel().cuda() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) accumulation_steps = 4 # 累积步数 global_step = 0 for epoch in range(num_epochs): for batch_idx, (inputs, labels) in enumerate(dataloader): inputs, labels = inputs.cuda(), labels.cuda() outputs = model(inputs) loss = loss_fn(outputs, labels) # 除以accumulation_steps以保持总损失尺度不变 loss = loss / accumulation_steps loss.backward() # 达到累积步数时更新参数 if (batch_idx + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad() global_step += 1

3.3 混合精度下的梯度累积

# 梯度累积 + 混合精度 scaler = GradScaler() accumulation_steps = 4 for batch_idx, (inputs, labels) in enumerate(dataloader): inputs, labels = inputs.cuda(), labels.cuda() with autocast(): outputs = model(inputs) loss = loss_fn(outputs, labels) / accumulation_steps scaler.scale(loss).backward() if (batch_idx + 1) % accumulation_steps == 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()

注意事项

  • Loss需要除以accumulation_steps，否则总损失会放大步数倍，导致学习率实际上升
  • Batch Normalization层的统计量在每个micro-batch上独立计算，batch size过小可能导致BN统计不稳定。可考虑使用SyncBN或GroupNorm替代
  • 梯度累积增加了通信和计算的同步开销，accumulation_steps不宜过大（通常2-8）
  • Learning Rate Warmup阶段建议使用较小的accumulation_steps或增大warmup步数

四、梯度裁剪 (Gradient Clipping)

4.1 梯度爆炸问题

在深度神经网络（尤其是RNN/LSTM、Transformer深层堆叠）训练中，梯度在反向传播过程中可能因连续矩阵乘法而指数级增长，导致梯度爆炸（gradient explosion）。梯度爆炸表现为loss突然变为NaN或Inf，参数更新幅度过大使模型彻底崩溃。梯度裁剪是最简单有效的预防手段。

4.2 两种裁剪方式

PyTorch提供了两种梯度裁剪函数，位于 torch.nn.utils 中：

按范数裁剪 (clip_grad_norm_)

# 按总范数裁剪（最常用） import torch.nn.utils as utils max_norm = 1.0 # 梯度的最大L2范数 # 计算所有参数的L2范数，若超过max_norm则等比例缩放到max_norm utils.clip_grad_norm_(model.parameters(), max_norm) optimizer.step()

clip_grad_norm_ 的工作原理：先计算所有参数梯度的全局L2范数 total_norm = sqrt(sum(grad_i^2))，若 total_norm > max_norm，则每个梯度乘以系数 max_norm / total_norm。这保持了梯度方向的正确性，仅调整步长。

按值域裁剪 (clip_grad_value_)

# 按值域裁剪 # 将每个梯度值限制在 [-clip_value, clip_value] 范围内 utils.clip_grad_value_(model.parameters(), clip_value=0.5) optimizer.step()

clip_grad_value_ 简单地将每个梯度元素裁剪到指定区间内。它改变了梯度的方向，因此更适用于某些对梯度方向不敏感的优化器，但通常不如按范数裁剪常用。

4.3 阈值选择策略

场景 推荐max_norm 说明
CNN图像分类 5.0 - 10.0 梯度爆炸风险较低，可设较大阈值
RNN/LSTM 0.25 - 1.0 长序列训练梯度爆炸风险高，需严格限制
Transformer 1.0 - 5.0 取决于层数和学习率
GAN 1.0 - 5.0 区分器和生成器的梯度尺度可能不同

诊断技巧

在训练早期记录梯度的总范数：

# 记录梯度范数用于调试 total_norm = 0.0 for p in model.parameters(): if p.grad is not None: param_norm = p.grad.data.norm(2) total_norm += param_norm.item() ** 2 total_norm = total_norm ** 0.5 print(f"Gradient norm: {total_norm:.4f}")

观察梯度范数在训练中的变化趋势，若经常接近设定阈值，则应增大阈值或排查网络结构是否存在数值不稳定因素。

4.4 与混合精度的配合

# 梯度裁剪 + 混合精度 + 梯度累积 scaler = GradScaler() accumulation_steps = 4 max_norm = 1.0 for batch_idx, (inputs, labels) in enumerate(dataloader): inputs, labels = inputs.cuda(), labels.cuda() with autocast(): outputs = model(inputs) loss = loss_fn(outputs, labels) / accumulation_steps scaler.scale(loss).backward() if (batch_idx + 1) % accumulation_steps == 0: # 先unscale梯度再裁剪 scaler.unscale_(optimizer) utils.clip_grad_norm_(model.parameters(), max_norm) scaler.step(optimizer) scaler.update() optimizer.zero_grad()

五、学习率调度 (Learning Rate Scheduling)

5.1 为什么需要学习率调度

学习率是深度学习中最重要的超参数之一。固定的学习率往往难以兼顾训练初期的快速收敛和后期的精细微调。学习率调度器（lr_scheduler）在训练过程中动态调整学习率，帮助模型跳出局部最优、加速收敛并提升最终性能。PyTorch的 torch.optim.lr_scheduler 提供了十余种调度策略。

5.2 常用调度器详解

StepLR：固定步长衰减

# 每隔 step_size 个epoch将LR乘以gamma scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=30, # 每30个epoch衰减一次 gamma=0.1 # 学习率乘以0.1 ) for epoch in range(num_epochs): train_one_epoch() scheduler.step() # 每个epoch后更新学习率 current_lr = scheduler.get_last_lr()[0]

CosineAnnealingLR：余弦退火

# 按余弦函数从初始LR下降到eta_min scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=50, # 半个周期长度（epoch数） eta_min=1e-6 # 最小学习率 ) for epoch in range(num_epochs): train_one_epoch() scheduler.step()

ReduceLROnPlateau：自适应降LR

# 当验证指标不再改善时降低学习率 scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', # 'min'监控loss下降，'max'监控accuracy上升 factor=0.5, # 新LR = 旧LR * factor patience=5, # 连续patience个epoch无改善则触发衰减 threshold=1e-4, # 判断改善的阈值 min_lr=1e-7, # LR下限 verbose=True # 打印LR变更日志 ) for epoch in range(num_epochs): train_loss = train_one_epoch() val_loss = validate() # 注意：ReduceLROnPlateau需要传入监控指标 scheduler.step(val_loss)

OneCycleLR：单周期策略

# 学习率先从低到高（warmup）再从高到低（anneal） scheduler = torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr=1e-3, # 峰值学习率 steps_per_epoch=len(dataloader), epochs=num_epochs, pct_start=0.3, # warmup占总步数的比例 anneal_strategy='cos' # 退火策略：'cos'或'linear' ) for epoch in range(num_epochs): for batch in dataloader: train_step(batch) scheduler.step() # 每个batch更新一次

5.3 调度器对比

调度器 适用场景 优势 劣势
StepLR 简单分类任务 直观、易调参 衰减时机不灵活
CosineAnnealingLR 图像分类、生成模型 平滑下降，常获更好结果 需配合warmup使用
ReduceLROnPlateau 验证指标易获取的任务 自适应、不需预设epoch 可能触发过晚
OneCycleLR 训练时间有限的场景 收敛极快、效果出色 超参数敏感

最佳实践

  • Warmup： 训练初期使用较小的学习率，逐步增加到目标值，可以避免模型在最开始的不稳定阶段发散。可结合 LinearLR 或 LambdaLR 实现
  • 组合调度： SequentialLR 可以按顺序组合多种调度策略
  • 监控LR曲线： 使用TensorBoard记录LR变化，确保LR衰减节奏与loss下降趋势匹配
  • 重启策略： CosineAnnealingWarmRestarts 定期重启LR，适用于逃离局部最优

5.4 组合调度实现

# Warmup + CosineAnnealing 组合 from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR warmup_epochs = 5 total_epochs = 100 warmup_scheduler = LinearLR( optimizer, start_factor=0.1, end_factor=1.0, total_iters=warmup_epochs ) cosine_scheduler = CosineAnnealingLR( optimizer, T_max=total_epochs - warmup_epochs, eta_min=1e-6 ) scheduler = SequentialLR( optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[warmup_epochs] ) for epoch in range(total_epochs): train_one_epoch() scheduler.step()

六、EMA指数移动平均 (Exponential Moving Average)

6.1 原理与动机

EMA (Exponential Moving Average) 的核心思想是维护一份模型参数的滑动平均值。在训练过程中，每个step更新原始模型参数后，以指数衰减的方式更新EMA参数副本：ema_param = decay * ema_param + (1 - decay) * model_param。由于训练后期的参数在局部最优点附近震荡，EMA平滑后的参数往往更接近损失函数的谷底中心，因此在验证和测试时使用EMA参数通常可获得更优的性能。

数学表达

设第 t 步的模型参数为 theta_t，EMA参数为 ema_t，衰减率为 alpha（通常取 0.999 或 0.9999）：
ema_0 = theta_0
ema_t = alpha * ema_{t-1} + (1 - alpha) * theta_t

在验证时使用 ema_t 替代 theta_t 进行前向传播。

6.2 标准实现

class EMA: """指数移动平均（Exponential Moving Average）""" def __init__(self, model, decay=0.999): self.model = model self.decay = decay self.shadow = {} self.backup = {} self.register() def register(self): for name, param in self.model.named_parameters(): if param.requires_grad: self.shadow[name] = param.data.clone() def update(self): for name, param in self.model.named_parameters(): if param.requires_grad: new_average = (self.decay * self.shadow[name] + (1.0 - self.decay) * param.data) self.shadow[name] = new_average.clone() def apply_shadow(self): """将EMA参数应用到模型（测试时使用）""" for name, param in self.model.named_parameters(): if param.requires_grad: self.backup[name] = param.data.clone() param.data.copy_(self.shadow[name]) def restore(self): """恢复原始模型参数（测试结束后调用）""" for name, param in self.model.named_parameters(): if param.requires_grad: param.data.copy_(self.backup[name]) self.backup = {} # 使用方式 ema = EMA(model, decay=0.999) for batch in dataloader: optimizer.zero_grad() loss = train_step(batch) loss.backward() optimizer.step() ema.update() # 每个step更新EMA参数 # 测试时应用EMA参数 ema.apply_shadow() test_accuracy = evaluate(model) ema.restore()

6.3 衰减率设置

decay的典型值与训练步数的关系：

训练总步数 推荐decay值 有效窗口
1k - 10k 0.99 ~100步
10k - 50k 0.999 ~1000步
50k+ 0.9999 ~10000步

实践建议

  • EMA在验证集上的表现通常比原始模型稳定，尤其在训练后期loss震荡时
  • 可将EMA作为模型集成的一种简易替代方案（实际上是对时间维度的集成）
  • 推荐在训练脚本中始终保存EMA参数，使得测试时可以选择使用
  • 若使用BN层，应用EMA参数后建议在验证集上重新前向传播以更新running statistics

七、分布式训练基础

7.1 DataParallel vs DistributedDataParallel

PyTorch提供了两种多GPU并行训练方案。torch.nn.DataParallel (DP) 是最简单的入门方式：它将输入batch分割到各GPU，每个GPU独立前向传播，主GPU收集梯度后更新参数。但DP存在严重性能瓶颈：主GPU需要承担所有梯度的归约（reduce）操作，通信开销随GPU数线性增长，且只支持单机多卡。

torch.nn.parallel.DistributedDataParallel (DDP) 是PyTorch推荐的方案。DDP采用多进程模型，每个GPU由一个独立进程控制，各进程加载模型副本并独立计算梯度，通过高效的Ring All-Reduce算法在后台异步同步梯度。DDP的通信开销与GPU数量呈亚线性关系，支持多机扩展，且不受Python GIL限制。

特性 DataParallel (DP) DistributedDataParallel (DDP)
进程模型 单进程多线程 多进程
通信方式 主GPU串行归约 Ring All-Reduce
性能 GPU增加时通信成为瓶颈 近线性扩展
多机支持 不支持 支持
适用场景 2-4 GPU快速实验 4+ GPU、多机训练

7.2 DDP标准实现

# train.py - 使用 torchrun 启动 # torchrun --nproc_per_node=4 train.py import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP import os def setup(): """初始化分布式进程组""" local_rank = int(os.environ['LOCAL_RANK']) world_size = int(os.environ['WORLD_SIZE']) dist.init_process_group( backend='nccl', init_method='env://', rank=int(os.environ['RANK']) ) torch.cuda.set_device(local_rank) return local_rank, world_size def cleanup(): dist.destroy_process_group() def main(): local_rank, world_size = setup() # 在每个GPU上创建独立的模型和数据加载器 model = MyModel().cuda(local_rank) model = DDP(model, device_ids=[local_rank]) # 使用DistributedSampler确保数据在GPU间不重复 from torch.utils.data.distributed import DistributedSampler sampler = DistributedSampler(dataset) dataloader = DataLoader(dataset, batch_size=64, sampler=sampler) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) for epoch in range(num_epochs): sampler.set_epoch(epoch) # 确保每个epoch数据不同顺序 for batch in dataloader: inputs, labels = batch inputs, labels = inputs.cuda(local_rank), labels.cuda(local_rank) outputs = model(inputs) loss = loss_fn(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() cleanup() if __name__ == "__main__": main()

7.3 启动方式

# 单机4卡启动 torchrun --nproc_per_node=4 train.py # 多机启动（node1: 2GPU, node2: 2GPU） # 在node1上： torchrun --nnodes=2 --nproc_per_node=2 \ --rdzv_endpoint=192.168.1.1:29500 \ --rdzv_id=my_job train.py # 在node2上（相同命令） torchrun --nnodes=2 --nproc_per_node=2 \ --rdzv_endpoint=192.168.1.1:29500 \ --rdzv_id=my_job train.py

DDP与混合精度结合要点

  • 每个进程独立创建GradScaler实例，缩放因子各自维护
  • autocast在每个GPU上独立运作，无需额外配置
  • 梯度裁剪前需要 scaler.unscale_(optimizer)，与单卡一致
  • 推荐使用 torchrun 启动而非手动 spawn，torchrun 自动处理环境变量设置和故障恢复

八、实验记录 (TensorBoard)

8.1 为什么需要实验记录

深度学习实验涉及大量超参数组合（学习率、batch size、网络层数、dropout比例等），仅靠终端输出无法有效追踪和分析实验过程。TensorBoard作为最广泛使用的可视化工具，可以记录标量指标（loss、accuracy、learning rate）、模型结构计算图、梯度/权重直方图、图像样本等多维度信息，帮助研究人员快速定位问题、比较实验效果。

8.2 基本用法

from torch.utils.tensorboard import SummaryWriter import numpy as np # 创建writer，指定日志目录 writer = SummaryWriter(log_dir='runs/mnist_experiment_01') # ======= 记录标量 ======= for step in range(num_steps): train_loss, train_acc = train_step() val_loss, val_acc = validate() writer.add_scalar('Loss/train', train_loss, step) writer.add_scalar('Loss/val', val_loss, step) writer.add_scalar('Accuracy/train', train_acc, step) writer.add_scalar('Accuracy/val', val_acc, step) writer.add_scalar('LearningRate', get_lr(optimizer), step) # ======= 记录直方图 ======= for name, param in model.named_parameters(): writer.add_histogram(f'Weights/{name}', param.data, step) if param.grad is not None: writer.add_histogram(f'Gradients/{name}', param.grad, step) # ======= 记录图像 ======= # images: [N, C, H, W] tensor, 值范围 [0, 1] writer.add_images('Input/Image', images, step) # ======= 记录模型结构图 ======= # 需要一次虚拟前向传播 dummy_input = torch.randn(1, 3, 224, 224).cuda() writer.add_graph(model, dummy_input) writer.close()

8.3 高级用法

# ======= 超参数记录与对比 ======= # 在实验结束时记录超参数和最终指标 hparams = { 'lr': 1e-3, 'batch_size': 64, 'optimizer': 'AdamW', 'dropout': 0.3, } metrics = { 'best_val_acc': best_acc, 'final_loss': final_loss, } writer.add_hparams(hparams, metrics) # ======= 嵌入向量可视化（t-SNE/PCA降维） ======= # 记录特征向量用于高维可视化 writer.add_embedding( features, # [N, D] 特征矩阵 metadata=labels, # [N] 类别标签 label_img=images, # [N, C, H, W] 原始图像（可选） global_step=step, tag='Features' ) # ======= 多实验对比 ======= # 在不同log_dir下创建writer，TensorBoard会自动叠加显示 writer1 = SummaryWriter('runs/exp_lr_1e-3') writer2 = SummaryWriter('runs/exp_lr_1e-4')

8.4 启动TensorBoard

# 启动TensorBoard（默认端口6006） tensorboard --logdir=runs --port=6006 # 浏览器访问 http://localhost:6006 # 多实验会自动显示在同一图表中，便于对比

记录策略建议

  • 标量记录不要太频繁：每N步记录一次即可，过于密集的记录会拖慢训练
  • 有意义的tag命名：使用 Category/Name 格式（如 Loss/train、Accuracy/val），TensorBoard会自动分组
  • 梯度直方图有助于检测梯度消失/爆炸：若某层权重直方图长时间不更新，说明梯度消失
  • 利用 add_hparams 建立超参数与性能的对应关系，便于调参复盘

九、torchinfo模型统计

9.1 为什么需要模型统计

在模型设计阶段，快速了解模型的参数量、计算量、各层输出尺寸等信息至关重要。传统的 print(model) 只能输出模块层次结构，无法直观展示每层的参数量和特征图尺寸变化。torchinfo（原 torchsummary 的升级版）提供了更完善的模型统计信息，包括每层输出shape、参数量、MACs（乘加操作数）等。

9.2 安装与基本使用

# 安装 pip install torchinfo
# 基本用法 from torchinfo import summary # 自动推断输入shape summary(model) # 或指定输入尺寸 summary(model, input_size=(32, 3, 224, 224)) # 更详细的配置 summary( model, input_size=(32, 3, 224, 224), # (batch, channels, height, width) batch_dim=0, # batch维度位置 col_names=[ # 显示的列 "input_size", "output_size", "num_params", "params_percent", "kernel_size", "mult_adds", ], depth=4, # 递归展开深度 verbose=1, # 输出详细程度 )

9.3 输出解读

# 示例输出（ResNet-18） ========================================================================================== Layer (type:depth-idx) Input Shape Output Shape Param # ========================================================================================== ResNet [32, 3, 224, 224] [32, 1000] -- ├─Conv2d: 1-1 [32, 3, 224, 224] [32, 64, 112, 112] 9,408 ├─BatchNorm2d: 1-2 [32, 64, 112, 112] [32, 64, 112, 112] 128 ├─ReLU: 1-3 [32, 64, 112, 112] [32, 64, 112, 112] -- ├─MaxPool2d: 1-4 [32, 64, 112, 112] [32, 64, 56, 56] -- ├─Sequential: 1-5 [32, 64, 56, 56] [32, 64, 56, 56] -- │ └─BasicBlock: 2-1 [32, 64, 56, 56] [32, 64, 56, 56] -- │ ├─Conv2d: 3-1 [32, 64, 56, 56] [32, 64, 56, 56] 36,864 │ ├─BatchNorm2d: 3-2 [32, 64, 56, 56] [32, 64, 56, 56] 128 │ ├─Conv2d: 3-3 [32, 64, 56, 56] [32, 64, 56, 56] 36,864 │ └─BatchNorm2d: 3-4 [32, 64, 56, 56] [32, 64, 56, 56] 128 ... ├─Linear: 1-6 [32, 512] [32, 1000] 513,000 ========================================================================================== Total params: 11,689,512 Trainable params: 11,689,512 Non-trainable params: 0 Total mult-adds (G): 1.82 ==========================================================================================

torchinfo vs torchsummary

  • torchinfo是torchsummary的社区维护升级版，支持更详细的信息列
  • 支持递归展开子模块（depth参数），便于分析复杂嵌套模型
  • 支持MACs（乘加操作数）统计，比FLOPs更接近实际硬件计算量
  • 自动处理多种输入类型（list、dict、tuple）
  • 经常与TensorBoard配合使用：先用torchinfo了解模型概况，再用TensorBoard深入分析训练过程
# 保存summary到文件以便查阅 import sys from io import StringIO buffer = StringIO() sys.stdout = buffer summary(model, input_size=(1, 3, 224, 224)) sys.stdout = sys.__stdout__ model_info = buffer.getvalue() with open('model_summary.txt', 'w') as f: f.write(model_info)

十、综合训练脚本示例

以下是一个整合了上述所有技术的完整训练循环骨架，展示了各项技术如何在实际项目中协同工作：

# comprehensive_train.py - 综合训练脚本 import torch import torch.nn as nn import torch.optim as optim import torch.nn.utils as utils from torch.cuda.amp import autocast, GradScaler from torch.utils.tensorboard import SummaryWriter from torchinfo import summary import os class Trainer: def __init__(self, model, device, config): self.model = model.to(device) self.device = device self.config = config # 优化器与调度器 self.optimizer = optim.AdamW( model.parameters(), lr=config['lr'] ) self.scheduler = optim.lr_scheduler.CosineAnnealingLR( self.optimizer, T_max=config['epochs'] ) # 混合精度 self.scaler = GradScaler(enabled=config['amp']) # 梯度累积步数 self.accum_steps = config.get('accum_steps', 1) # EMA（可选） self.ema = None if config.get('ema_decay'): self.ema = EMA(model, decay=config['ema_decay']) # TensorBoard log_dir = config.get('log_dir', 'runs/exp') self.writer = SummaryWriter(log_dir=log_dir) # 打印模型结构 summary(model, input_size=config['input_size']) def train_epoch(self, dataloader, epoch): self.model.train() total_loss = 0.0 self.optimizer.zero_grad() for batch_idx, (inputs, labels) in enumerate(dataloader): inputs, labels = inputs.to(self.device), labels.to(self.device) # 混合精度前向传播 with autocast(enabled=self.config['amp']): outputs = self.model(inputs) loss = nn.CrossEntropyLoss()(outputs, labels) loss = loss / self.accum_steps # 反向传播（带GradScaler） self.scaler.scale(loss).backward() # 达到累积步数时更新参数 if (batch_idx + 1) % self.accum_steps == 0: # 梯度裁剪（先unscale） if self.config.get('max_norm'): self.scaler.unscale_(self.optimizer) utils.clip_grad_norm_( self.model.parameters(), self.config['max_norm'] ) self.scaler.step(self.optimizer) self.scaler.update() self.optimizer.zero_grad() # 更新EMA if self.ema: self.ema.update() # 记录标量 global_step = epoch * len(dataloader) + batch_idx self.writer.add_scalar( 'Loss/train', loss.item() * self.accum_steps, global_step ) total_loss += loss.item() * self.accum_steps return total_loss / len(dataloader) def validate(self, dataloader): """验证时使用EMA参数（如果可用）""" if self.ema: self.ema.apply_shadow() self.model.eval() correct = total = 0 with torch.no_grad(): for inputs, labels in dataloader: inputs, labels = inputs.to(self.device), labels.to(self.device) outputs = self.model(inputs) preds = outputs.argmax(dim=1) correct += (preds == labels).sum().item() total += labels.size(0) accuracy = correct / total if self.ema: self.ema.restore() return accuracy def fit(self, train_loader, val_loader, epochs): for epoch in range(epochs): train_loss = self.train_epoch(train_loader, epoch) val_acc = self.validate(val_loader) self.scheduler.step() # 记录验证指标和LR self.writer.add_scalar('Accuracy/val', val_acc, epoch) self.writer.add_scalar( 'LearningRate', self.scheduler.get_last_lr()[0], epoch ) print(f"Epoch {epoch+1}/{epochs} | " f"Train Loss: {train_loss:.4f} | " f"Val Acc: {val_acc:.4f}") self.writer.close()

十一、核心要点总结

十二、进一步思考与实践

进阶学习路径

  1. DeepSpeed/FSDP： 理解ZeRO优化器的三个阶段，掌握大规模模型（如LLaMA、GPT）的分布式训练方案
  2. 自动混合精度（AMP）原理： 深入理解FP16/BF16/FP32的数值表示差异，掌握FP16梯度下溢的根本原因
  3. 梯度检查点（Gradient Checkpointing）： 以时间换空间，在长序列Transformer训练中极其实用
  4. 学习率重启动（Cosine Annealing with Warm Restarts）： 通过周期性提升学习率帮助模型跳出局部最优
  5. SWA/SWAG（Stochastic Weight Averaging）： EMA的进阶版本，在训练末期对参数做均匀平均，进一步平滑损失景观
  6. WandB/MLflow： 替代TensorBoard的云端实验管理平台，支持团队协作和自动化超参数搜索

推荐资源

  • PyTorch官方文档：pytorch.org/docs/stable/amp.html
  • NVIDIA混合精度训练白皮书：docs.nvidia.com/deeplearning/performance/mixed-precision-training
  • DDP官方教程：pytorch.org/tutorials/intermediate/ddp_tutorial.html
  • TensorBoard文档：tensorboard.dev