1. 从一行 import 说起:PyTorch 到底在解决什么问题
很多人第一次接触 PyTorch,是从import torch这一行开始的。但真正让 PyTorch 在众多深度学习框架里站稳脚跟的,不是它“能跑神经网络”这件事本身——能跑神经网络的框架太多了——而是它对计算过程的表达方式。我见过不少从其他框架迁移过来的开发者,最初都会有一个共同的困惑:为什么 PyTorch 的代码看起来这么“像普通的 Python 程序”?答案恰恰就在这里。
PyTorch 的核心定位是一个基于动态计算图的深度学习框架。这句话里每个词都值得拆开看。“动态计算图”意味着计算流程是在代码运行时逐行构建的,而不是像静态图那样先定义完整结构再执行。这个差异带来的直接好处是:你可以用 Python 的if、for、while去控制网络的前向过程,可以在调试器里逐行跟踪张量的变化,可以随时打印中间结果。对于做研究、做实验、做原型验证的人来说,这种“所见即所得”的体验几乎是决定性的。
那它适合谁?我的判断是三类人。第一类是刚入门深度学习的学生和转行者,因为 PyTorch 的 API 设计贴近 NumPy 的直觉,学习曲线相对平缓。第二类是做算法研究和快速迭代的工程师,动态图带来的灵活性在实验阶段价值极高。第三类是需要把模型部署到生产环境的开发者,因为 PyTorch 这几年在torch.compile、TorchScript、ONNX 导出、量化等方向持续投入,从研究到落地的链路已经相当完整。
这篇文章不打算写成一份 API 手册——那种东西官方文档已经做得很好了。我想做的是把 PyTorch 使用过程中那些文档里不会重点讲、但实际项目中一定会遇到的东西讲清楚:张量和自动求导的底层逻辑、动态图的真实代价、训练循环里那些看似不起眼却影响巨大的细节、显存管理的实战技巧、以及从实验代码走向可复现工程时踩过的坑。无论你是刚跑通第一个 MNIST 的新手,还是已经写过几万行训练代码的老手,希望都能从中找到一点有用的东西。
2. 张量与自动求导:PyTorch 的两块地基
2.1 张量不只是“多维数组”那么简单
从表面看,torch.Tensor就是一个多维数组,和 NumPy 的ndarray很像。但两者的区别远不止“能不能放到 GPU 上”。PyTorch 张量身上挂着一整套元信息:dtype、device、layout、requires_grad、grad_fn,这些属性决定了它在计算图中的角色。
先说dtype。深度学习里最常见的坑之一就是类型不匹配。比如你用torch.tensor([1, 2, 3])创建了一个整型张量,然后想把它和一个浮点模型参数相乘,PyTorch 会直接报错而不是默默帮你转换。这其实是好事——NumPy 那种隐式类型提升在调试时经常让人抓狂。我的习惯是:凡是参与梯度计算的张量,一律显式指定dtype=torch.float32,不要依赖默认推断。
再说device。CPU 和 GPU 之间的张量不能直接运算,这是新手最常撞的墙。一个典型的报错是Expected all tensors to be on the same device。解决办法很简单,但关键是养成习惯:在训练脚本开头定义一个device = torch.device('cuda' if torch.cuda.is_available() else 'cpu'),然后所有张量创建和模型.to(device)都统一走这个变量。我见过太多项目里到处硬编码.cuda(),结果换到没有 GPU 的机器上就全线崩溃。
关于张量的创建,有几个效率上的细节值得注意。torch.tensor(data)会拷贝数据,而torch.as_tensor(data)在可能的情况下会共享内存。如果你从 NumPy 数组转过来且不打算修改原数据,用as_tensor更省内存。另外,torch.zeros、torch.ones、torch.empty这几个函数里,empty不初始化内存,速度最快,但里面的值是随机的——如果你打算立刻覆盖它,用empty是合理的;如果忘了覆盖,就会得到一堆看似正常实则垃圾的数值,这种 bug 极难排查。
2.2 自动求导机制:requires_grad与计算图的构建
PyTorch 的自动求导(autograd)是整个框架最精妙的部分。理解它的关键在于搞清楚计算图是什么时候、以什么方式被构建出来的。
当你对一个requires_grad=True的张量做任何运算,PyTorch 都会在后台记录这个运算,形成一个有向无环图。每个结果张量都有一个grad_fn属性,指向创建它的那个运算节点。比如:
import torch x = torch.tensor([2.0, 3.0], requires_grad=True) y = x ** 2 z = y.sum() print(z.grad_fn) # <SumBackward0 object>这里的z.grad_fn就是计算图的入口。调用z.backward()时,PyTorch 会沿着这个图反向传播,把梯度累加到每个叶子张量的.grad属性上。注意“累加”这个词——PyTorch 的梯度是累加的,不是覆盖的。这意味着如果你在训练循环里忘了调用optimizer.zero_grad(),梯度会不断累积,模型更新就会完全错乱。这是新手最常见的 bug 之一,而且症状很隐蔽:loss 可能一开始还在下降,然后突然爆炸或者停滞。
还有一个容易忽略的点:只有叶子节点的.grad会被保留。中间结果的梯度在反向传播后默认会被释放。如果你想查看中间层的梯度,需要用retain_grad()。这个设计是为了节省内存,但在调试时经常让人困惑“为什么我的.grad是 None”。
关于torch.no_grad()的使用,也值得单独说。在推理阶段或者任何不需要计算梯度的场合,一定要用with torch.no_grad():包起来。原因不只是“省去计算梯度的开销”,更重要的是它会阻止计算图的构建,从而大幅降低显存占用。我在早期项目里做过对比:同一个模型做推理,不加no_grad时显存占用是加了之后的将近两倍。对于大模型来说,这个差异可能就是“能跑”和“OOM”的区别。
2.3 动态图的代价:灵活性的另一面
动态图带来的灵活性是 PyTorch 的招牌,但它不是没有代价的。每次前向传播都要重新构建计算图,这意味着框架无法在运行前对计算做全局优化。静态图框架可以在编译期做算子融合、内存复用规划等优化,而动态图只能边跑边决定。
PyTorch 对这个问题的回应是torch.compile(在 2.0 版本引入)。它的思路是用一个编译器在运行时捕获计算图,然后做优化并生成更高效的代码。实测下来,在 Transformer 类模型上通常能有 20% 到 50% 的加速,具体取决于模型结构和输入形状是否稳定。但torch.compile也不是万能的:如果模型里有大量依赖 Python 控制流的逻辑,或者输入形状频繁变化,编译带来的收益可能被反复编译的开销抵消。
我的经验是:实验阶段用 eager 模式(默认模式)保证灵活性,确定模型结构后再尝试torch.compile做加速。不要一上来就 compile,否则调试会变得很痛苦,因为编译后的代码报错信息往往不如 eager 模式直观。
3. 训练循环里那些决定成败的细节
3.1 一个标准训练循环的解剖
几乎每个 PyTorch 项目的核心都是一个训练循环。它的骨架看起来很简单:
for epoch in range(num_epochs): model.train() for batch in dataloader: inputs, labels = batch inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() model.eval() with torch.no_grad(): # 验证逻辑 pass但就是这十几行代码,藏着大量容易出错的细节。我按执行顺序逐个说。
model.train()和model.eval()这一对调用,影响的是Dropout和BatchNorm这类层的行为。Dropout在训练时随机置零,在推理时不做;BatchNorm在训练时用当前 batch 的统计量,在推理时用running mean和running variance。如果你忘了切换模式,模型在验证集上的表现会莫名其妙地差,而且这种错误不会报错,只会让你怀疑人生。我的习惯是把model.train()写在 epoch 循环内部而不是外部,虽然效果一样,但这样更不容易漏掉。
optimizer.zero_grad()的位置也有讲究。它必须放在loss.backward()之前。有些人喜欢放在optimizer.step()之后,效果一样,但放在 backward 之前更符合直觉。另外,PyTorch 1.7 之后zero_grad有一个set_to_none参数,默认是False,意味着它把梯度置零而不是释放。设置成True可以省一点显存,但会让依赖.grad非 None 的代码出问题。一般项目里保持默认就好。
loss.backward()之后、optimizer.step()之前,是梯度裁剪的黄金位置。RNN、Transformer 这类模型很容易出现梯度爆炸,加一行torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)往往能救命。这个操作的原理是把所有参数的梯度看作一个向量,如果它的范数超过max_norm,就按比例缩放。max_norm取多少?1.0 是一个常见的默认值,但具体要看模型和任务,我一般会在 0.5 到 5.0 之间试。
3.2 DataLoader 的并行陷阱与调优
DataLoader是数据管道的核心,它的num_workers参数决定了用几个子进程加载数据。设成 0 表示在主进程里加载,设成大于 0 的值会启动子进程。这里有几个坑。
第一个坑是Windows 上的多进程问题。在 Windows 上,num_workers > 0要求你的训练代码放在if __name__ == '__main__':保护块里,否则会无限递归地启动子进程。这个坑在 Linux 上不存在,所以很多从 Linux 迁移到 Windows 的代码会突然崩溃。
第二个坑是num_workers不是越大越好。每个 worker 都会占用内存,而且进程间通信本身有开销。我的经验是:num_workers设成 CPU 核心数的一半到全部之间,具体要看数据加载的复杂度和 batch size。如果数据加载本身很快(比如数据已经在内存里),设大了反而拖慢速度。可以用time.time()简单测一下一个 epoch 的耗时,然后调整。
第三个坑是worker 里的随机性。如果你在Dataset的__getitem__里用了随机数做数据增强,每个 worker 的随机种子默认是不同的,这会导致每个 epoch 的增强结果不可复现。解决办法是在DataLoader里设置worker_init_fn,给每个 worker 分配确定的种子。这个细节在写论文或者做需要严格复现的实验时特别重要。
还有一个实用技巧:pin_memory=True。当使用 GPU 时,这个选项会让 DataLoader 把数据放在锁页内存里,从而加速 CPU 到 GPU 的传输。实测下来,对于小 batch 或者数据传输频繁的场景,提升很明显。配合non_blocking=True使用效果更好。
3.3 学习率调度:不只是“调小一点”
学习率是训练中最重要的超参数,没有之一。但很多人对学习率调度的理解停留在“训练后期调小一点”。实际上,调度策略的选择对最终性能影响很大。
最常用的是余弦退火(Cosine Annealing)。它的思路是让学习率沿着余弦曲线从初始值平滑下降到接近零。相比阶梯式下降,余弦退火的优势在于下降过程平滑,不会在阶梯处造成 loss 的突变。PyTorch 里用torch.optim.lr_scheduler.CosineAnnealingLR实现。
另一个常用的是带热重启的余弦退火(Cosine Annealing with Warm Restarts)。它周期性地把学习率重置回初始值,让模型有机会跳出局部最优。这个策略在集成学习里特别有用——你可以用不同的重启周期训练多个模型然后集成。
还有一个容易被忽视的是预热(Warmup)。在训练最开始,模型参数是随机的,此时用大学习率容易导致训练不稳定。预热就是在最初几百到几千步里把学习率从很小的值线性增加到初始值。Transformer 类模型几乎都需要预热,这也是原始 Transformer 论文里的标准做法。
需要特别注意的是:PyTorch 1.1 之后,学习率调度器的调用时机变了。在旧版本里,scheduler.step()应该放在 epoch 循环里;在新版本里,如果用的是ReduceLROnPlateau之外的调度器,通常也放在 epoch 级别。但如果你用的是按 step 调度的策略,就要放在 batch 循环里。这个细节搞错的话,学习率曲线会完全不对。我的建议是:每次用调度器时,打印几次学习率看看是否符合预期,这个习惯能帮你避免很多隐蔽的问题。
4. 显存管理:从“能跑”到“跑得稳”
4.1 显存都去哪了:拆解模型的内存占用
很多人第一次遇到CUDA out of memory时,第一反应是“模型太大了”。但实际上,模型参数本身占的显存往往只是一小部分。一个完整的训练过程,显存主要被四部分占用:
| 占用来源 | 说明 | 大致比例 |
|---|---|---|
| 模型参数 | 权重和偏置 | 1x |
| 梯度 | 每个参数对应一个梯度 | 1x |
| 优化器状态 | 如 Adam 的动量和方差 | 2x(Adam) |
| 中间激活值 | 前向传播中保存的中间结果 | 视模型而定,可能远大于前三者 |
对于 Adam 优化器,参数、梯度、优化器状态加起来就是参数量的 4 倍。一个 1 亿参数的模型,用 float32 存储,光这部分就是 400MB 乘以 4 等于 1.6GB。而中间激活值在大 batch 或长序列场景下可能轻松超过这个数。
理解了这一点,优化显存就有了方向。如果瓶颈在优化器状态,可以换用 SGD 或者 Adafactor 这类省显存的优化器;如果瓶颈在激活值,可以用梯度检查点(Gradient Checkpointing)——它用计算换显存,在前向时不保存中间激活,反向时重新计算。PyTorch 里用torch.utils.checkpoint.checkpoint实现。
4.2 混合精度训练:省显存又提速
混合精度训练(Mixed Precision Training)是这几年最实用的显存优化手段之一。它的核心思想是:前向和反向传播用 float16(或 bfloat16),参数更新用 float32。float16 占的显存是 float32 的一半,而且现代 GPU 对 float16 运算有专门的加速硬件。
PyTorch 从 1.6 开始内置了torch.cuda.amp,用起来很简单:
from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() for inputs, labels in dataloader: optimizer.zero_grad() with autocast(): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这里的GradScaler是为了解决 float16 的数值范围问题。float16 能表示的最小正数大约是 6e-8,很多小梯度会下溢成零。GradScaler的做法是先把 loss 放大一个系数,反向传播得到放大后的梯度,更新前再缩回去。这样小梯度就不会丢失。
实测下来,混合精度通常能省 30% 到 50% 的显存,同时提速 20% 到 100%,具体取决于模型和 GPU 架构。但要注意:不是所有操作都支持 float16。一些自定义的 CUDA 算子或者特殊的损失函数可能在 float16 下出问题。如果遇到 NaN 或者结果异常,可以先关掉混合精度排查。
4.3 梯度累积:小显存跑大 batch
当你想要大 batch 但显存不够时,梯度累积是一个经典方案。它的原理很简单:多次前向反向累积梯度,再统一更新一次参数。这样等效于大 batch,但每次只处理小 batch。
accumulation_steps = 4 for i, (inputs, labels) in enumerate(dataloader): outputs = model(inputs) loss = criterion(outputs, labels) / accumulation_steps loss.backward() if (i + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()注意 loss 要除以accumulation_steps,否则梯度会放大对应倍数。另外,如果你用了 BatchNorm,梯度累积会改变实际的 batch 统计量——因为每个小 batch 的统计量是独立的。这种情况下可以考虑用 GroupNorm 或者 SyncBatchNorm 替代。
梯度累积和混合精度可以叠加使用,两者不冲突。我在显存受限的机器上经常这么组合,效果很稳。
5. 从实验代码到可复现工程:那些绕不开的坑
5.1 随机种子:复现性的第一道门槛
“为什么我跑出来的结果和上次不一样?”这是深度学习里最常见的疑问之一。答案几乎总是:随机性没有被完全控制。
PyTorch 的随机性来源有好几个:参数初始化、Dropout、数据打乱、数据增强、以及 CUDA 的一些非确定性算子。要完全固定这些,需要一套组合拳:
import torch import numpy as np import random 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 = Falsecudnn.deterministic = True会让 cuDNN 只用确定性的算法,代价是可能慢一些。cudnn.benchmark = False则关闭了算法自动选择,进一步保证确定性。
但即使这样,完全复现仍然很难。多 GPU 训练、某些自定义算子、以及不同版本的 PyTorch 都可能引入差异。我的建议是:不要追求 bit-level 的完全复现,而是追求统计意义上的可复现——同样的代码和种子,跑出来的指标在合理误差范围内一致,这就够了。
5.2 模型保存与加载:state_dict才是正道
保存模型时,新手容易犯的错误是直接torch.save(model, path),把整个模型对象序列化。这样做的问题很多:模型类必须能被 pickle,代码重构后可能加载失败,而且文件体积大。
正确的做法是保存state_dict:
torch.save(model.state_dict(), 'model.pth') # 加载 model = MyModel() model.load_state_dict(torch.load('model.pth')) model.to(device)如果要保存训练状态以便断点续训,还需要保存优化器状态、epoch、学习率调度器状态等:
checkpoint = { 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), 'loss': loss, } torch.save(checkpoint, 'checkpoint.pth')加载时注意torch.load的map_location参数。如果你在 GPU 上保存的模型要在 CPU 上加载,不加map_location='cpu'会报错。反过来也一样。这个细节在跨机器迁移时经常被忽略。
5.3 分布式训练:从 DataParallel 到 DDP
当单卡显存不够或者训练太慢时,就需要多卡。PyTorch 提供了两种主要方案:DataParallel(DP)和DistributedDataParallel(DDP)。
DataParallel用起来最简单,一行model = nn.DataParallel(model)就行。但它的性能很差,因为每个 forward 都要把模型复制到所有 GPU,而且主 GPU 的负载和显存占用远高于其他卡。我的建议是:除非是快速验证,否则不要用 DataParallel。
DistributedDataParallel是官方推荐的方案。它每个进程控制一张卡,通过 NCCL 后端做梯度同步,效率高得多。但配置起来复杂一些,需要设置init_process_group、DistributedSampler、local_rank等。一个最小的 DDP 训练脚本大概长这样:
import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backend='nccl') local_rank = int(os.environ['LOCAL_RANK']) torch.cuda.set_device(local_rank) model = MyModel().to(local_rank) model = DDP(model, device_ids=[local_rank]) sampler = DistributedSampler(dataset) dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)DDP 里有一个容易忽略的点:每个 epoch 开始时要调用sampler.set_epoch(epoch),否则每个 epoch 的数据划分顺序都一样,会损害训练效果。这个细节在官方教程里提了,但很多人会漏掉。
另外,DDP 下的 BatchNorm 默认是各卡独立的,如果 batch size 很小,统计量会不准。这时可以用SyncBatchNorm把统计量跨卡同步,代价是通信开销增加。
6. 踩坑实录:几个让我印象深刻的排查过程
6.1 loss 变成 NaN:从表象到根因的完整链路
有一次训练一个 Transformer 模型,前几百步 loss 正常下降,然后突然变成 NaN。这种问题最烦人,因为它不是一开始就错,而是跑了一段时间才出问题。
我的排查链路是这样的:
第一步,确认 NaN 出现的位置。我在 loss 计算后加了一个检查:if torch.isnan(loss): print('NaN detected'),然后逐步往前找,发现是某一层的输出先出现了 NaN。
第二步,检查学习率。学习率过大是导致 NaN 的常见原因。我把学习率调小了一个数量级,NaN 出现的时间推迟了,但没有消失。这说明学习率不是唯一原因。
第三步,检查数据。我用torch.isnan(inputs).any()检查输入数据,发现某些 batch 里确实有 NaN。追查数据源,发现是数据预处理时某个除法操作遇到了零。修复数据后,问题解决。
这个经历给我的教训是:NaN 问题要从数据、模型、优化器三个方向同时排查,不要只盯着一个。而且,加一个数据校验的步骤(比如在 Dataset 里检查 NaN 和 Inf)能省下大量调试时间。
6.2 显存泄漏:为什么训练越跑越慢
另一个印象深刻的坑是显存泄漏。现象是:训练刚开始时显存占用正常,跑了几百个 batch 后显存逐渐涨满,最后 OOM。
排查这类问题的关键是区分“显存占用高”和“显存泄漏”。前者是静态的,后者是动态增长的。我用torch.cuda.memory_allocated()在每个 batch 后打印显存占用,确认了确实是持续增长。
原因出在一个不起眼的地方:我在验证阶段把一些中间结果存进了一个列表,用于后续分析,但忘了在验证结束后清空。这个列表持有张量的引用,导致这些张量的显存无法释放。解决办法很简单:验证结束后del掉列表,或者用.detach().cpu()把需要的数据转到 CPU 上。
这个坑的教训是:任何长期持有的张量引用都可能导致显存泄漏。在训练循环里,除了模型参数和优化器状态,其他张量都应该在不需要时及时释放。用del加torch.cuda.empty_cache()可以强制回收,但empty_cache本身有开销,不要频繁调用。
6.3 多卡训练速度反而变慢:通信开销的真相
有一次我把单卡训练改成 4 卡 DDP,预期是速度提升接近 4 倍,结果只提升了不到 2 倍,有时候甚至比单卡还慢。
排查后发现两个问题。第一个是数据加载成了瓶颈。单卡时num_workers=4够用,4 卡时每个进程还是 4 个 worker,总共 16 个 worker 抢 CPU,反而导致数据供给跟不上。把每个进程的num_workers降到 2 后,情况改善。
第二个问题是梯度同步的通信开销。模型参数量大时,每次反向传播后的 all-reduce 操作会占用大量时间。解决办法是用梯度分桶(gradient bucketing)——DDP 默认会把梯度分成多个桶,桶越大通信效率越高,但显存占用也越大。可以通过bucket_cap_mb参数调整。另外,如果网络带宽有限,可以考虑用梯度压缩或者更高效的通信后端。
这个经历让我明白:多卡训练不是简单的“卡越多越快”,数据管道、通信带宽、模型大小都会影响实际加速比。上线多卡之前,一定要做小规模的 benchmark,确认瓶颈在哪里。
7. 一些让代码更顺手的实战习惯
7.1 用torchinfo替代print(model)
print(model)只能看到模型的层级结构,看不到每层的输出形状和参数量。torchinfo(以前叫torchsummary)能给出更完整的信息:
from torchinfo import summary summary(model, input_size=(batch_size, channels, height, width))它会输出每层的输入输出形状、参数量、是否可训练等信息。在调试模型结构时,这个工具能帮你快速发现维度不匹配的问题。
7.2 用torch.profiler找性能瓶颈
当训练速度不如预期时,不要靠猜,用 profiler:
from torch.profiler import profile, ProfilerActivity with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: # 训练几步 pass print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))它会告诉你时间花在了哪些算子上。我经常用它发现一些意外的情况,比如某个数据预处理操作占用了大量 CPU 时间,或者某个自定义算子效率极低。
7.3 用torch.utils.tensorboard记录训练过程
TensorBoard 是可视化训练过程的标准工具。PyTorch 内置了SummaryWriter:
from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter('runs/experiment_1') writer.add_scalar('Loss/train', loss, epoch) writer.add_scalar('Loss/val', val_loss, epoch) writer.add_scalar('LR', optimizer.param_groups[0]['lr'], epoch)把 loss、学习率、梯度范数、权重分布都记录下来,训练出问题时回看这些曲线,往往能快速定位。我特别推荐记录梯度范数——它能提前预警梯度爆炸或消失。
7.4 配置文件管理:别把超参数硬编码
小项目里把超参数写在代码里没问题,但项目一大,改一个参数就要翻遍代码。用配置文件(YAML 或 JSON)管理超参数是更好的做法:
# config.yaml model: hidden_size: 768 num_layers: 12 training: batch_size: 32 learning_rate: 1e-4 epochs: 50然后用argparse或omegaconf加载。这样实验记录和复现都方便得多。我习惯把每次实验的配置文件连同 checkpoint 一起保存,这样几个月后回头看还能知道当时用了什么参数。
8. 写在最后:PyTorch 教会我的事
用了这么多年 PyTorch,我最大的体会是:框架的便利性会掩盖底层的复杂性,但复杂性不会消失,只会在出问题时集中爆发。动态图让你写代码很爽,但理解计算图的构建和释放机制,才能在显存爆炸时知道从哪下手;自动求导让你不用手写反向传播,但理解梯度累加和叶子节点的概念,才能避免那些“看起来能跑但结果不对”的 bug。
另一个体会是:不要迷信默认值。PyTorch 的默认参数是为了通用性设计的,不是为你的具体任务优化的。num_workers、pin_memory、学习率、batch size、梯度裁剪阈值,这些都需要根据实际情况调整。调整的依据不是“别人说这样好”,而是你自己的 profiling 和实验数据。
最后,PyTorch 的生态还在快速演进。torch.compile、torch.export、torchao这些新工具正在改变从训练到部署的流程。保持学习,保持实验,但也要记住:工具是为人服务的,理解问题本身永远比记住 API 更重要。