pytorch-patterns

用途

PyTorch 开发模式（PyTorch Development Patterns）是一套惯用的 PyTorch 模式与最佳实践，用于构建稳健、高效且可复现的深度学习应用。它适用于以下场景：编写新的 PyTorch 模型或训练脚本、审查深度学习代码、调试训练循环或数据管道、优化 GPU 显存占用或训练速度、搭建可复现的实验环境。

核心原则

设备无关代码：始终编写同时适用于 CPU 和 GPU 的代码，不硬编码设备。推荐使用 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")，再通过 .to(device) 放置模型与数据；避免直接调用 .cuda()，否则在没有 GPU 时会崩溃。
可复现性优先：设置所有随机种子以获得可复现结果。示例函数 set_seed 会设置 torch.manual_seed、torch.cuda.manual_seed_all、np.random.seed、random.seed，并将 torch.backends.cudnn.deterministic 设为 True、torch.backends.cudnn.benchmark 设为 False。
显式形状管理：始终记录并验证张量形状。在 forward 中通过注释标注每一步的形状变化，例如输入 (batch_size, channels, height, width)，经过卷积、池化、view 展平后输出 (batch_size, num_classes)。

模型架构模式

清晰的 nn.Module 结构：将特征提取与分类器分别组织为 nn.Sequential，在 __init__ 中定义层，在 forward 中只做前向计算；避免在 forward 中动态创建权重。
正确的权重初始化：通过 model.apply(model._init_weights) 对 Linear、Conv2d、BatchNorm2d 等模块显式初始化，例如对 Linear/Conv2d 使用 nn.init.kaiming_normal_，对 BatchNorm2d 使用 nn.init.ones_ 与 nn.init.zeros_。

训练循环模式

标准训练循环：训练前调用 model.train()；使用 optimizer.zero_grad(set_to_none=True) 更高效地清梯度；支持混合精度训练（torch.amp.autocast 与 torch.amp.GradScaler）；使用 torch.nn.utils.clip_grad_norm_ 裁剪梯度。
验证循环：使用 @torch.no_grad() 装饰器，并始终调用 model.eval()，以禁用 dropout 并使用 BatchNorm 的运行统计量。

数据管道模式

自定义 Dataset：实现 __len__ 与 __getitem__，使用类型提示，在 __getitem__ 中打开图像、转换 RGB、应用 transform 并返回张量与标签。
高效 DataLoader 配置：设置 shuffle=True、num_workers=4、pin_memory=True、persistent_workers=True、drop_last=True，以加速数据加载与 CPU 到 GPU 的传输，并保持 BatchNorm 的批次大小一致。
变长数据的自定义 collate：在 collate_fn 中使用 nn.utils.rnn.pad_sequence 将序列填充到批次内最大长度。

检查点模式

保存检查点时应包含 epoch、model_state_dict、optimizer_state_dict 与 loss，以便恢复训练；加载时使用 torch.load(path, map_location="cpu", weights_only=True)，再分别加载模型与优化器状态。仅保存模型权重会导致无法恢复训练。

性能优化

混合精度训练：使用 torch.amp.GradScaler("cuda") 与 torch.amp.autocast("cuda")。
大模型的梯度检查点：使用 torch.utils.checkpoint.checkpoint，以计算换显存，示例中传入 use_reentrant=False。
torch.compile 加速：PyTorch 2.0+ 可使用 torch.compile(model, mode="reduce-overhead")，模式包括 "default"（安全）、"reduce-overhead"（更快）、"max-autotune"（最快）。

快速参考：PyTorch 惯用法

model.train() / model.eval()：训练/评估前始终设置模式
torch.no_grad()：推理时禁用梯度
optimizer.zero_grad(set_to_none=True)：更高效的梯度清除
.to(device)：设备无关的张量/模型放置
torch.amp.autocast：混合精度，约 2 倍速度
pin_memory=True：更快的 CPU→GPU 数据传输
torch.compile：JIT 编译加速（2.0+）
weights_only=True：更安全的模型加载
torch.manual_seed：可复现实验
gradient_checkpointing：以计算换显存

应避免的反模式

验证时忘记 model.eval()：会导致 dropout 仍生效、BatchNorm 使用批次统计量。
破坏自动求导的原地操作：如 F.relu(x, inplace=True) 或 x += residual；应使用非原地操作。
在训练循环内反复将模型移到 GPU：应在循环前一次性 model.to(device)。
在 backward 前调用 .item()：会从计算图中分离，导致无法反向传播；.item() 只应在反向传播后用于日志。
未正确使用 torch.save：直接保存整个模型脆弱且不可移植，应保存 state_dict。

限制

资料未提供该技能的安装命令、依赖版本、平台支持范围、收费方式或具体使用步骤，这些信息暂未提供。