写代码这件事做到后面拼的往往是平时的积累。就拿 PyTorch 来说很多功能官方文档里写得明明白白但真到写项目的时候谁还会去翻文档找 API大家都习惯直接翻自己以前的代码或者去网上搜一段能用的片段。这篇文章就是干这个用的——把我这些年写 PyTorch 时沉淀下来的常用代码段整理成一份可以直接抄的合集涉及数据加载、模型定义、训练循环、断点续训、迁移学习这些高频场景。每种场景我都会说明这段代码解决了什么问题、为什么要这么写、哪些地方容易踩坑。不管你是刚入门 PyTorch 的初学者还是写了一阵子但总觉得代码不够顺畅的开发者这份合集应该都能帮你省下不少查资料的功夫。1. 数据加载与预处理让 DataLoader 真正为你干活1.1 Dataset 的自定义写法与常见误区数据加载是所有训练流程的地基。虽然 PyTorch 官方提供了很多现成的 Dataset但实际项目里大部分情况还是要自己写一个。自定义 Dataset 的核心就两个方法__len__和__getitem__。前者返回样本总数后者根据索引返回一条样本通常是(输入, 标签)的元组。from torch.utils.data import Dataset class MyDataset(Dataset): def __init__(self, data_list, transformNone): self.data_list data_list self.transform transform def __len__(self): return len(self.data_list) def __getitem__(self, idx): # 这里根据你的数据格式返回样本 sample self.data_list[idx] # 可以在这里做索引、裁剪、归一化等操作 if self.transform: sample self.transform(sample) return sample, label # 返回输入和标签这里最容易被忽略的是__getitem__里到底该不该做预处理。我的习惯是只做轻量级的、跟数据本身绑定的操作比如读图、归一化、转 Tensor。而像随机裁剪、翻转这类数据增强操作交给transform对象去管。这样 Dataset 更纯粹也方便以后更换增强策略。很多人会在 Dataset 里直接做标准化然后发现训练效果不稳定。原因可能是标准化参数均值和方差是在整批数据上计算的而 Dataset 应该只负责单样本读取。标准化放在 transform 里或者放训练脚本里算好后再传进来这样逻辑更清晰也方便调试。1.2 Transform 组合别把顺序写反了PyTorch 的torchvision.transforms提供了一系列图像处理操作用Compose把它们串起来。顺序很讲究常见的合理顺序是from torchvision import transforms transform_train transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) transform_val transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])有几个细节值得说明一下RandomCrop之前最好先Resize到一个稍大的尺寸否则crop出来的区域信息量不够。验证集和测试集不要用随机增强只需要Resize到最终大小。Normalize的参数是 ImageNet 的均值和标准差如果你做的是自研数据集最好自己统计一下数据集的真实均值和方差用 ImageNet 的参数可能是次优选择。ToTensor()会把 HWC 的 uint8 数组转成 CHW 的 float32数值从 [0,255] 缩放到 [0,1]。所以Normalize一定要放在ToTensor后面否则你归一化的是 [0,255] 的值效果完全不对。这些点单看都不难但组合在一起尤其容易乱。一个简单的自查方法训练集用了哪些随机变换验证集就要去掉这些随机变换训练集和测试集的预处理必须保持一致。1.3 DataLoader 参数num_workers和pin_memory是白送的加速from torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue ) val_loader DataLoader( val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue, drop_lastFalse )num_workers表示用几个子进程加载数据。设为 0 表示在主进程里加载数据读取会成为训练瓶颈设大了又可能因为进程间通信开销反而变慢。一般建议从 4 开始调如果 CPU 核心够多可以继续加大但超过 8 之后收益很小。pin_memoryTrue在显卡训练时能把数据传输到 GPU 的效率明显提高这个参数几乎总是建议打开。drop_last在训练集设为 True 可以避免最后一个不完整的 batch 对 BatchNorm 统计量造成扰动。2. 模型定义与初始化网络结构之外的隐藏细节2.1 权重初始化不初始化等于在裸奔PyTorch 的nn.Conv2d、nn.Linear等层自带默认初始化多数情况下可以直接训练。但如果你用了自定义层或者想让收敛更稳定手动初始化就很有必要了。做法通常是在模型定义完参数之后对需要的层分别设置初始化方式。import torch.nn as nn def init_weights(m): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 0.01) nn.init.constant_(m.bias, 0) model.apply(init_weights)model.apply()是沿着模块树递归调用init_weights的所以不管嵌套多少层都能覆盖到。这段代码里fan_out是配合 ReLU 激活函数的常见选择因为这样能保证前向传播时每一层输出的方差变化不大。如果换成了 Tanh 或 Sigmoid初始化策略可能要做相应调整。2.2 forward 里的一个隐藏陷阱原地操作class MyNet(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 64, 3, padding1) self.relu nn.ReLU(inplaceTrue) def forward(self, x): x self.conv1(x) x self.relu(x) return xinplaceTrue的意思是直接在输入张量上做修改省一份内存。但要注意inplace之后这个张量不能再参与需要记录反向传播的操作。最典型的就是x F.relu(x) # OK y x.view(x.size(0), -1) # OK x x 1 # 如果之前是 inplace 的 relu这里也 OK但如果先提取了x的某个中间值然后对这个中间值做 inplace 操作反向传播就会出错。一个稳妥的习惯是不明确需要省内存时ReLU、dropout这些就用默认的inplaceFalse省得排查半天。3. 训练循环、梯度裁剪与学习率调度训练主循环的规范写法3.1 一个标准训练循环背后的细节几乎所有 PyTorch 训练的代码都长一个样但细看各家写法有些细节的处理并不一样。我自己习惯的模板是import torch import torch.nn as nn def train_one_epoch(model, train_loader, optimizer, criterion, device, clip_gradNone): model.train() total_loss 0 total_correct 0 total_num 0 for batch_idx, (inputs, targets) in enumerate(train_loader): inputs, targets inputs.to(device), targets.to(device) # 前向传播 outputs model(inputs) loss criterion(outputs, targets) # 梯度清零、反向传播、参数更新 optimizer.zero_grad() loss.backward() if clip_grad is not None: torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad) optimizer.step() # 统计 total_loss loss.item() * inputs.size(0) pred outputs.argmax(dim1) total_correct (pred targets).sum().item() total_num inputs.size(0) if batch_idx % 50 0: print(f batch {batch_idx}: loss{loss.item():.4f}) return total_loss / total_num, total_correct / total_num这里有几个细节loss.item()会打断计算图取的是一个 Python 浮点数如果不取.item()而是直接做total_loss loss计算图会一直不释放内存越积越多训练到后面就 OOM 了。这个坑我踩过排查了很久才发现是统计这行的问题。optimizer.zero_grad()必须在loss.backward()之前调用否则梯度会累加到上一轮的梯度上。现代版本 PyTorch 会在optimizer.step()之后自动清零部分状态的但手动写zero_grad()是明确且可靠的做法。model.train()放在循环最外层调用确保 BatchNorm 和 Dropout 进入训练模式。3.2 梯度裁剪防止梯度爆炸的最简单防线梯度爆炸在 RNN、Transformer 这类模型里非常常见特征就是 loss 突然变成 NaN 或者很大。clip_grad_norm_的常用阈值是 1.0 到 5.0设得太小会影响训练速度设得太大就起不到防爆炸的作用。torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)它的原理是计算所有参数梯度的总范数如果超过max_norm就按比例缩放。跟clip_grad_value_不一样clip_grad_value_是直接把每个梯度值裁剪到一个固定区间后者更暴力但会失真更严重一些。我一般先用clip_grad_norm_1.0起步如果训练稳定再尝试增大。3.3 学习率调度器别用一套固定策略从训练到结束学习率可以说是训练里最需要动态调整的超参数。PyTorch 提供了多种调度策略最常见的有几种# 每隔 step_size 个 epoch 衰减一个因子 scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1) # 当指标如验证集 loss不再下降时降低学习率 scheduler torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemin, factor0.5, patience5) # 余弦退火配合 warm restart 使用 scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50)ReduceLROnPlateau很有意思它不属于LR调度基类的标准调用方式而是需要你在每个 epoch 验证完之后手动scheduler.step(val_loss)。它的好处是不用预设固定的衰减节点而是根据实际验证指标动态调整非常适合不知道训练到什么时候 loss 会平台化的场景。3.4 EMA 滑动平均低成本提升模型稳定的常用技巧EMAExponential Moving Average在很多比赛中被证明能让模型更稳定代码也不复杂。它的思路是维护一组“平均参数”每个 step 用当前参数和历史上的平均参数做加权平均推理时用平均参数而不是原始参数。def ema_update(model, ema_model, decay0.999): with torch.no_grad(): for p, ema_p in zip(model.parameters(), ema_model.parameters()): ema_p.data.mul_(decay).add_(p.data, alpha1 - decay)这个decay越大历史信息保留越多模型越平滑。一般取 0.99 到 0.999 之间具体要看训练 step 总量。注意这里的ema_model要在一开始就把model的参数复制一份过去并在整个训练过程中只做更新、不做梯度计算。推理时加载ema_model的权重即可。这里的开销几乎可以忽略但对最终精度的收益往往很明显。4. 随机种子与断点续训两个“关键时刻救命”的代码段4.1 随机种子固定从数据加载到 cuDNN 全套设置训练可复现性是很多实验的硬要求。PyTorch 里固定随机种子要比想象中涉及的层面多光设torch.manual_seed是不够的。import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) # 尽量保证卷积等操作的确定性 torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False关于benchmark多说一句benchmarkTrue会让 cuDNN 在第一次前向时自动选择最合适的卷积算法从而加速训练但它每次选择算法时有一定的随机性可能导致复现不稳定。如果要严格复现就设benchmarkFalse如果只追求速度性能优先优先可以保留benchmarkTrue。这个不是 bug而是复现性和性能之间的一种取舍。另外即使种子固定了如果训练过程中用到了多进程数据加载num_workers的子进程仍然可能存在非确定性。一个补充办法是给每个 worker 单独设置种子def worker_init_fn(worker_id): np.random.seed(42 worker_id) train_loader DataLoader(..., num_workers4, worker_init_fnworker_init_fn)这种做法能保证每个 worker 的数据采样序列也确定下来尤其是在做随机增强时会有效果。4.2 断点续训不只是保存模型参数训练过程意外中断是最让人崩溃的事尤其训练了几天的模型一下子前功尽弃。断点续训的核心是保存足够多的状态包括模型参数、优化器状态、学习率调度器状态、当前 epoch 数以及一些额外的指标。def save_checkpoint(model, optimizer, scheduler, epoch, val_acc, save_path): checkpoint { model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict(), epoch: epoch, val_acc: val_acc, best_val_acc: best_val_acc, } torch.save(checkpoint, save_path)加载时要注意把模型先构造好再依次加载各个状态def load_checkpoint(model, optimizer, scheduler, checkpoint_path, device): checkpoint torch.load(checkpoint_path, map_locationdevice) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) scheduler.load_state_dict(checkpoint[scheduler_state_dict]) start_epoch checkpoint[epoch] 1 val_acc checkpoint.get(val_acc, 0.0) return model, optimizer, scheduler, start_epoch, val_accmap_locationdevice这个参数特别重要。如果你是在 GPU 上训练的然后换到 CPU 机器上加载不写这个参数会直接报错。反过来也一样。省内存的做法还有很多比如只保存model.state_dict()而不是整个 model 对象因为torch.save(model)会把网络结构、类定义位置等信息一并序列化跨环境部署时很容易出问题。这也是我建议用state_dict保存模型的原因。4.3 保存最优模型而不是最后一个模型正常训练里我们会同时维护一个“当前最优模型”用验证集的指标来判断是否更新。写法也不难best_acc 0.0 for epoch in range(start_epoch, total_epochs): train_loss, train_acc train_one_epoch(...) val_acc evaluate(model, val_loader, device) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth)这里有个很容易忽略的逻辑best_acc必须在加载断点时一并恢复否则断点续训之后best_acc又变成 0期间保存的模型权重无法正确被覆盖更新。5. 模型保存加载、迁移学习与微调工程落地的习惯写法5.1 多 GPU 训练后的模型加载时注意module.前缀很多教程里都推荐优先用nn.DataParallel来处理多 GPU 训练写法是一行代码model nn.DataParallel(model)但这个包装会改变state_dict的 key每个参数名都会多出module.前缀。如果你用单卡训练完一个模型再想在多卡环境跑推理或者反过来直接加载可能会报尺寸不匹配的错。解决方式有两种。第一种是把加载的权重去掉module.前缀state_dict torch.load(model.pth) from collections import OrderedDict new_state_dict OrderedDict() for k, v in state_dict.items(): name k[7:] if k.startswith(module.) else k new_state_dict[name] v model.load_state_dict(new_state_dict)第二种是在加载之前又把模型包了一层 DataParallelkey 正好能对上。实际使用中建议保存模型权重时保存最原始的model.state_dict()而不要保存nn.DataParallel包装后的state_dict这样后续不论单卡多卡都方便加载。5.2 迁移学习冻结特征层只训练分类头训练数据不够时迁移学习几乎是标配做法。常见的套路是加载预训练模型的权重把最后一层替换成自己数据集的分类数其它层要么冻结、要么用很小的学习率微调。import torchvision.models as models # 加载预训练模型 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_classes 10 # 替换分类头 in_features model.fc.in_features model.fc nn.Linear(in_features, num_classes) # 冻结除 fc 外的所有层 for name, param in model.named_parameters(): if fc not in name: param.requires_grad False optimizer torch.optim.Adam(model.fc.parameters(), lr1e-3)这里要注意的是冻结后依然会有 BatchNorm 层的问题。即使requires_gradFalse前向传播时 BatchNorm 仍然会更新它内部的running_mean和running_var这两个不是可学习的参数。所以冻结特征层后如果训练数据量很少最好用model.eval()模式跑特征提取或者尝试把 BatchNorm 也冻结只放开分类头的训练否则统计量可能被不稳定的数据扰动。更进一步的做法是分阶段微调先冻结 backbone 只训练新分类头跑几个 epoch 之后解锁部分层用更小的学习率继续微调。这个策略在数据有限时比一次性全量微调稳定得多。5.3 推理时的随机性控制torch.no_grad 与推理模式推理阶段我们不需要反向传播PyTorch 提供了两种上下文管理器torch.no_grad()和torch.inference_mode()。二者的差异在于推理模式还会进一步禁用一些自动求导的功能速度能快一点点。实际使用中我会在评估函数外面包上torch.inference_mode()def evaluate(model, val_loader, device): model.eval() correct 0 total 0 with torch.inference_mode(): for inputs, targets in val_loader: inputs, targets inputs.to(device), targets.to(device) outputs model(inputs) pred outputs.argmax(dim1) correct (pred targets).sum().item() total targets.size(0) return correct / totalmodel.eval()只改变模式标志让 BatchNorm 和 Dropout 使用推理行为但它本身不关闭梯度计算。所以必须搭配no_grad或inference_mode才能真正关闭自动求导减少显存占用。6. 写完模型之后自查清单我踩过的几个真实深坑6.1 忘记 model.train() 和 model.eval() 切换很多人训练循环里只写model.train()忘记在验证前切换到评估模式。比如 BatchNorm 层的running_mean和running_var只有在训练模式下才会更新如果验证和训练混用一套模式很有可能会导致验证指标忽高忽低、难以判断真实效果。正确习惯是在每个 epoch 训练开始前model.train()验证开始前model.eval()两个模式都要写别偷懒。6.2 数据没有同步到同一个 device当输入数据在 CPU、模型在 GPU 时直接前向会报 RuntimeError。有的框架会自动帮你搬数据但 PyTorch 不会。写训练代码时建议统一在最前面inputs, targets inputs.to(device), targets.to(device)。如果你看到Expected tensor for argument #1 input to have the same device as tensor for weight就是这个原因。6.3 BatchNorm 在 batch_size 太小时的统计失真BatchNorm 在训练时用的是每个 batch 内部的均值和方差如果 batch_size 很小比如只有 2 或 4统计量会有比较大的噪声容易导致训练不稳定。如果模型里有大量 BatchNorm 层而你的显存又只能装下很小的 batch可以考虑几个方案增大 batch 的替代做法是使用梯度累积见下一小节、换成 GroupNorm、或者用同步 BatchNorm。但批量比较小的时候GroupNorm 往往更稳定。6.4 梯度累积显存不够时的换时间方案有些场景下 batch_size64 才符合惯例但显存只能跑 batch_size16这时可以用梯度累积模拟大 batchaccumulation_steps 4 optimizer.zero_grad() for batch_idx, (inputs, targets) in enumerate(train_loader): inputs, targets inputs.to(device), targets.to(device) outputs model(inputs) loss criterion(outputs, targets) loss loss / accumulation_steps # 归一化避免梯度总体过大 loss.backward() if (batch_idx 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()这里有个细节值得说明loss / accumulation_steps这步很多人会忽略。因为梯度求的是 loss 对参数的偏导4 次小 batch 的 loss 加起来再做一次优化等效于合并的一个大 batch。如果不除 accumulation_steps等于梯度莫名放大了 4 倍跟直接用 batch_size64 还是有区别的。6.5 尽早保存最佳模型别等训练跑完再保存训练时间比较长的时候中途有可能掉点比如学习率过大、数据扰动造成 loss 发散如果只留存最后一个 epoch 的模型很可能不是整个训练过程中验证指标最好的版本。所以我强烈建议在每个 epoch 结束之后都执行一次验证然后判断是否刷新best_acc若是就保存一次模型。在最后再针对 best 的 checkpoint 做一次推理评估通常结果会优于 epoch 末直接用 last 模型。一点个人整理习惯最后分享一个我自己的习惯。因为平时项目多、代码碎片也多我会按“数据加载-模型定义-训练循环-工具函数”这四个层级来维护一套自己的模板文件。每次开新项目先复制模板再按项目需求改动。改的时候只动参数不动结构等积累了新的有效写法再回填到模板里。这样带来的好处是每个项目都自然地继承了之前沉淀的最佳实践踩过的坑也不会因为换了项目就忘掉。如果你也是从零开始整理自己的工具代码不用一上来就追求大全。先把你现在项目里最需要的那部分写好比如自定义 Dataset 和训练循环跑通一次之后再去补其它模块。代码段不在于多在于每段都清楚自己解决什么问题。真到用的时候能直接拿来抄、抄了不出错这才是“真香”这两个字的含义。