企业官网建设流程全解析

写代码这件事，做到后面拼的往往是平时的积累。就拿 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, transform=None): 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(brightness=0.2, contrast=0.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_size=32, shuffle=True, num_workers=4, pin_memory=True, drop_last=True ) val_loader = DataLoader( val_dataset, batch_size=32, shuffle=False, num_workers=4, pin_memory=True, drop_last=False )

num_workers表示用几个子进程加载数据。设为 0 表示在主进程里加载，数据读取会成为训练瓶颈；设大了又可能因为进程间通信开销反而变慢。一般建议从 4 开始调，如果 CPU 核心够多可以继续加大，但超过 8 之后收益很小。pin_memory=True在显卡训练时能把数据传输到 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, mode='fan_out', nonlinearity='relu') 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, padding=1) self.relu = nn.ReLU(inplace=True) def forward(self, x): x = self.conv1(x) x = self.relu(x) return x

inplace=True的意思是直接在输入张量上做修改，省一份内存。但要注意，inplace之后这个张量不能再参与需要记录反向传播的操作。最典型的就是：

x = F.relu(x) # OK y = x.view(x.size(0), -1) # OK x = x + 1 # 如果之前是 inplace 的 relu，这里也 OK

但如果先提取了x的某个中间值，然后对这个中间值做 inplace 操作，反向传播就会出错。一个稳妥的习惯是：不明确需要省内存时，ReLU、dropout这些就用默认的inplace=False，省得排查半天。

3. 训练循环、梯度裁剪与学习率调度：训练主循环的规范写法

3.1 一个标准训练循环背后的细节

几乎所有 PyTorch 训练的代码都长一个样，但细看各家写法，有些细节的处理并不一样。我自己习惯的模板是：

import torch import torch.nn as nn def train_one_epoch(model, train_loader, optimizer, criterion, device, clip_grad=None): 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(dim=1) 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_norm=1.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_size=10, gamma=0.1) # 当指标（如验证集 loss）不再下降时，降低学习率 scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5) # 余弦退火，配合 warm restart 使用 scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)

ReduceLROnPlateau很有意思，它不属于LR调度基类的标准调用方式，而是需要你在每个 epoch 验证完之后手动scheduler.step(val_loss)。它的好处是不用预设固定的衰减节点，而是根据实际验证指标动态调整，非常适合不知道训练到什么时候 loss 会平台化的场景。

3.4 EMA 滑动平均：低成本提升模型稳定的常用技巧

EMA（Exponential Moving Average）在很多比赛中被证明能让模型更稳定，代码也不复杂。它的思路是维护一组“平均参数”，每个 step 用当前参数和历史上的平均参数做加权平均，推理时用平均参数而不是原始参数。

def ema_update(model, ema_model, decay=0.999): with torch.no_grad(): for p, ema_p in zip(model.parameters(), ema_model.parameters()): ema_p.data.mul_(decay).add_(p.data, alpha=1 - 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(seed=42): 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多说一句：benchmark=True会让 cuDNN 在第一次前向时自动选择最合适的卷积算法，从而加速训练，但它每次选择算法时有一定的随机性，可能导致复现不稳定。如果要严格复现就设benchmark=False；如果只追求速度，性能优先优先可以保留benchmark=True。这个不是 bug，而是复现性和性能之间的一种取舍。

另外，即使种子固定了，如果训练过程中用到了多进程数据加载，num_workers的子进程仍然可能存在非确定性。一个补充办法是给每个 worker 单独设置种子：

def worker_init_fn(worker_id): np.random.seed(42 + worker_id) train_loader = DataLoader(..., num_workers=4, worker_init_fn=worker_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_location=device) 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_acc

map_location=device这个参数特别重要。如果你是在 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)

第二种是在加载之前又把模型包了一层 DataParallel，key 正好能对上。实际使用中，建议保存模型权重时保存最原始的model.state_dict()而不要保存nn.DataParallel包装后的state_dict，这样后续不论单卡多卡都方便加载。

5.2 迁移学习：冻结特征层，只训练分类头

训练数据不够时，迁移学习几乎是标配做法。常见的套路是加载预训练模型的权重，把最后一层替换成自己数据集的分类数，其它层要么冻结、要么用很小的学习率微调。

import torchvision.models as models # 加载预训练模型 model = models.resnet18(weights=models.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(), lr=1e-3)

这里要注意的是，冻结后依然会有 BatchNorm 层的问题。即使requires_grad=False，前向传播时 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(dim=1) correct += (pred == targets).sum().item() total += targets.size(0) return correct / total

model.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_size=64 才符合惯例，但显存只能跑 batch_size=16，这时可以用梯度累积模拟大 batch：

accumulation_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_size=64 还是有区别的。

6.5 尽早保存最佳模型，别等训练跑完再保存

训练时间比较长的时候，中途有可能掉点（比如学习率过大、数据扰动造成 loss 发散），如果只留存最后一个 epoch 的模型，很可能不是整个训练过程中验证指标最好的版本。所以我强烈建议，在每个 epoch 结束之后都执行一次验证，然后判断是否刷新best_acc，若是就保存一次模型。在最后再针对 best 的 checkpoint 做一次推理评估，通常结果会优于 epoch 末直接用 last 模型。

一点个人整理习惯

最后分享一个我自己的习惯。因为平时项目多、代码碎片也多，我会按“数据加载-模型定义-训练循环-工具函数”这四个层级来维护一套自己的模板文件。每次开新项目，先复制模板，再按项目需求改动。改的时候只动参数，不动结构；等积累了新的有效写法，再回填到模板里。这样带来的好处是，每个项目都自然地继承了之前沉淀的最佳实践，踩过的坑也不会因为换了项目就忘掉。

如果你也是从零开始整理自己的工具代码，不用一上来就追求大全。先把你现在项目里最需要的那部分写好，比如自定义 Dataset 和训练循环，跑通一次之后再去补其它模块。代码段不在于多，在于每段都清楚自己解决什么问题。真到用的时候能直接拿来抄、抄了不出错，这才是“真香”这两个字的含义。