PyTorch实战指南：从动态图原理到生产部署的完整链路

1. 为什么一个“pytorch”标题值得单独写一篇

你如果是在技术社区里闲逛，看到有人只甩出一个“pytorch”当标题，第一反应大概率是：这哥们要么是懒得写标题，要么是准备写一篇大杂烩。但我自己从最早用0.3版本一路踩坑到现在，越来越觉得这个词本身就是一个巨大的话题入口。它不是一个框架名那么简单，它代表了一整套从研究到生产的思维方式的转变。很多人第一次接触深度学习，就是被import torch这行代码带进门的。它解决了什么问题？简单说，它让“定义网络、求梯度、跑在GPU上”这三件原本需要大量样板代码的事情，变得像写普通Python一样自然。适合谁看？如果你正在犹豫要不要从别的框架转过来，或者你已经会写模型但总觉得训练过程像黑盒，又或者你是个刚入门的新手，被各种教程里五花八门的写法搞晕了，那这篇东西就是给你准备的。我不会只给你列API，我会把我在实际项目里怎么选、怎么调、怎么排错的经验，全部拆开揉碎讲清楚。核心关键词就一个：pytorch，但围绕它展开的，是动态图、自动求导、张量计算、GPU加速、模型部署这一整条链路。你把这篇文章当成一个老手在咖啡馆里跟你聊他这些年用pytorch的心得就行，没有教科书式的定义，只有能直接抄作业的步骤和踩过的坑。

2. 核心设计思路：为什么是动态图，为什么是Python优先

2.1 动态计算图到底解决了什么痛点

我最早用的是静态图那一套，先定义好整个计算图，再喂数据进去跑。那种感觉就像你先画好一张极其复杂的电路图，然后才能通电。好处是图一旦建好，优化空间大，部署也方便。但坏处太明显了：你想在训练过程中根据中间结果动态调整网络结构，比如处理变长序列或者实现一些复杂的控制流，静态图就会让你写得非常别扭。pytorch的动态图机制，官方叫define-by-run，意思是计算图是在你执行代码的过程中实时构建的。你写一个for循环，每次迭代都根据当前张量的值决定下一步做什么，这在pytorch里是天然支持的。我举个生活化的类比：静态图就像你出门前必须把一整天的路线全部规划好，哪怕路上遇到修路也不能改；动态图就像你边走边看导航，随时可以拐弯。对于研究人员来说，这种灵活性是致命的吸引力。你可以在forward函数里直接写if-else，直接打印中间张量的值，直接用Python的调试器打断点。这大大缩短了从想法到验证的周期。我实测下来，用动态图调试一个自定义的注意力机制，比在静态图框架里省了至少一半的时间，因为你不需要去理解那些复杂的图构建API，你只需要会Python。

2.2 Python原生体验带来的连锁反应

pytorch另一个让我回不去的地方，是它和Python生态的无缝融合。你可以用numpy的思维去操作张量，可以用matplotlib直接画图，可以用pdb调试，可以用列表推导式去处理数据。这种“原生感”带来的效率提升是巨大的。我见过很多新手，在别的框架里被各种占位符、会话、图作用域搞得晕头转向，转到pytorch之后第一反应是：“原来可以这么写？”比如定义一个简单的全连接网络，你只需要继承nn.Module，然后在__init__里声明层，在forward里写计算逻辑。没有tf.placeholder，没有tf.Session，没有tf.graph。这种设计哲学的背后，是pytorch团队对开发者体验的极致追求。他们赌的是：研究人员更在乎灵活性和调试便利性，而不是极致的静态优化。事实证明这个赌注赢了。现在你去翻最新的论文，附带的代码十有八九是pytorch写的。这种生态效应又反过来让pytorch的社区越来越活跃，你遇到问题几乎都能搜到答案。我个人的体会是，当你不需要花精力去对抗框架本身的时候，你才能把精力真正放在模型设计和数据上。

2.3 从研究到生产的桥梁设计

很多人有个误解，觉得pytorch只适合研究，部署不行。这个印象在早期确实存在，但现在已经完全过时了。pytorch的设计思路里，从一开始就埋下了从研究平滑过渡到生产的伏笔。最典型的就是torch.jit和torch.export这套工具链。你可以在研究阶段用动态图快速迭代，等模型结构稳定了，再用torch.jit.trace或者torch.jit.script把模型转换成静态图表示，然后导出成TorchScript，脱离Python环境运行。这个过程不需要你重写模型，只需要加几行代码。我做过一个图像分类的项目，研究阶段用动态图调了两个月，最后部署到C++环境里，转换过程只花了一个下午。当然中间也踩了一些坑，比如有些Python原生语法TorchScript不支持，需要改写，但整体迁移成本远低于我的预期。另外还有ONNX导出这条路，可以把模型转到其他推理引擎上。这种“研究时怎么爽怎么来，部署时再收紧”的思路，我觉得是非常务实的。它不像某些框架那样，从一开始就要求你按照部署的标准来写代码，那样会扼杀很多实验性的想法。

3. 核心细节解析：张量、自动求导与模块化设计

3.1 张量操作里那些容易翻车的细节

张量是pytorch里最基础的数据结构，你可以把它理解成一个可以跑在GPU上的多维数组。但就是这最基础的东西，我见过太多人栽跟头。第一个坑是view和reshape的区别。view要求张量在内存里是连续的，如果不连续就会报错，而reshape会自动处理连续性，返回一个新的张量或者视图。我刚开始用的时候，经常在permute或者transpose之后直接调view，然后程序就崩了。后来养成了习惯：不确定的时候就用reshape，确定内存布局的时候再用view。第二个坑是原地操作。像add_、mul_这种带下划线的操作是原地修改，在需要计算梯度的场景下要特别小心，因为自动求导系统可能会因为原地修改而丢失历史信息。我一般只在明确不需要梯度的预处理阶段用原地操作，比如数据归一化。第三个坑是设备管理。你经常会在CPU和GPU之间来回搬数据，如果忘了.to(device)，就会报设备不匹配的错误。我的经验是，在代码开头就定义好device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')，然后在创建张量和模型的时候都显式指定设备。还有一个细节是数据类型，默认的浮点类型是torch.float32，但有些操作比如索引要求整数类型，混用会报错。我习惯在数据加载的时候就把类型统一好，避免在模型里做隐式转换。

3.2 自动求导机制背后的逻辑与实操禁忌

自动求导是pytorch的灵魂。你只需要在张量上设置requires_grad=True，然后所有基于这个张量的运算都会被记录，最后调用.backward()就能自动算出梯度。听起来很简单，但里面有几个关键点必须搞清楚。第一，梯度是累加的，不是覆盖的。这意味着你在每次迭代开始前必须调用optimizer.zero_grad()或者手动把梯度置零，否则梯度会越加越大，训练直接发散。我刚开始就犯过这个错误，loss不降反升，排查了半天才发现是忘了清梯度。第二，torch.no_grad()的使用场景。在验证和推理阶段，你不需要计算梯度，这时候用with torch.no_grad():包裹起来，可以显著减少显存占用和计算时间。我实测过，在验证集上加上这个上下文管理器，显存占用能降三分之一左右。第三，detach()的作用。有时候你想把某个张量从计算图中剥离出来，比如在强化学习里计算目标Q值的时候，就需要用detach()切断梯度回传。这个操作不会改变张量的值，只是让它不再追踪历史。第四，关于梯度裁剪。当你的网络比较深或者用了RNN结构时，梯度爆炸是很常见的问题。我一般会在backward()之后、step()之前加上torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)，这个操作简单但极其有效，能避免很多训练不稳定的情况。

3.3 nn.Module的模块化哲学与常见误用

nn.Module是pytorch里组织模型的基石。你定义的所有层、整个网络，都继承自这个类。它的核心机制是自动注册参数和子模块。当你在__init__里把一个nn.Linear赋值给self.fc时，这个层的权重和偏置就自动注册到了整个网络的参数列表里，你调用model.parameters()就能拿到所有需要优化的参数。这个设计非常优雅，但有几个常见的误用。第一个误用是在__init__里用普通的Python列表来存层，比如self.layers = [nn.Linear(10, 10) for _ in range(3)]。这样写这些层的参数不会被注册，优化器看不到它们，训练的时候这些层根本不会更新。正确的做法是用nn.ModuleList或者nn.Sequential。第二个误用是在forward里创建新的层。每次前向传播都创建新层，不仅效率低，而且参数无法复用。所有的层都应该在__init__里定义好。第三个是关于train()和eval()模式。有些层在训练和推理时的行为不一样，比如Dropout和BatchNorm。你必须记得在训练循环里调用model.train()，在验证时调用model.eval()。我见过有人忘了切换，导致验证结果波动很大，排查了很久才发现是Dropout还在起作用。还有一个细节是，即使你调用了model.eval()，如果忘了用torch.no_grad()，梯度还是会计算，只是Dropout和BatchNorm的行为变了。这两个要配合使用。

4. 实操过程：从零搭建一个完整的训练流水线

4.1 数据准备与Dataset、DataLoader的配合

任何模型训练都始于数据。pytorch提供了Dataset和DataLoader两个抽象来帮你处理数据加载。Dataset是一个抽象类，你需要实现__len__和__getitem__两个方法。__len__返回数据集大小，__getitem__根据索引返回一个样本。我一般会把数据预处理也放在__getitem__里，比如图像增强、归一化、张量转换。这样做的好处是预处理可以并行化，因为DataLoader可以开多个工作进程。DataLoader负责批处理、打乱、多进程加载。这里有几个参数需要特别注意。batch_size不用多说，根据显存大小来定。shuffle在训练时设为True，验证时设为False。num_workers控制加载数据的子进程数，我一般设为CPU核心数的一半左右，设太大反而会因为进程间通信开销导致速度下降。pin_memory设为True可以加速CPU到GPU的数据传输，但前提是你用的是GPU。还有一个我踩过的坑：在Windows系统上，num_workers大于0时，如果__getitem__里用了lambda或者局部函数，可能会报pickle错误。解决办法是把所有辅助函数定义在模块顶层，或者把num_workers设为0。另外，如果你自定义了collate_fn来处理变长序列，记得在里面用pad_sequence做填充，并且要处理好填充部分的mask，否则会影响loss计算。

4.2 模型定义、损失函数与优化器的选择逻辑

模型定义我习惯遵循一个模板：继承nn.Module，在__init__里声明所有层，在forward里定义计算流。对于简单的全连接网络，用nn.Sequential就够了。对于复杂的结构，比如多输入多输出，就需要自己写forward。损失函数的选择取决于任务。分类任务最常用的是nn.CrossEntropyLoss，它内部已经包含了softmax，所以你的模型输出不需要再加softmax层。回归任务用nn.MSELoss或者nn.L1Loss。如果遇到类别不平衡，可以用nn.CrossEntropyLoss(weight=class_weights)来给少数类更高的权重。优化器的选择上，Adam是我最常用的默认选项，因为它对学习率不那么敏感，收敛也快。但如果你追求极致的泛化性能，SGD加上动量往往能取得更好的结果，只是需要更精细的学习率调度。我一般会先用Adam快速跑一个baseline，确认模型结构没问题，然后再换SGD调参。学习率调度器也很重要，StepLR、CosineAnnealingLR、ReduceLROnPlateau我都用过。ReduceLROnPlateau比较省心，它根据验证集loss自动降低学习率，不需要你手动设定下降的时机。但要注意它的patience参数，设得太小会导致学习率降得太快，模型还没收敛就停了。

4.3 训练循环的骨架与关键监控指标

一个标准的训练循环包含以下几个步骤：前向传播、计算损失、反向传播、更新参数。我习惯把训练和验证写成一个函数，接受模型、数据加载器、优化器、损失函数等参数。在训练循环里，我会记录每个epoch的平均loss和准确率。准确率的计算要注意，如果是多分类，需要用torch.argmax(outputs, dim=1)拿到预测类别，然后和标签比较。这里有个细节：在计算准确率的时候，记得用.detach()把张量从计算图中分离出来，避免不必要的显存占用。我还会用tqdm来显示进度条，这样能直观看到每个batch的处理速度。监控指标方面，除了loss和准确率，我还会关注学习率的变化、梯度范数、以及显存占用。梯度范数可以用torch.nn.utils.clip_grad_norm_的返回值来获取，如果发现梯度范数突然变得很大，说明可能出现了梯度爆炸，需要检查数据或者调整网络结构。显存占用可以用torch.cuda.memory_allocated()来查看，如果接近显存上限，就需要减小batch size或者优化模型。我一般会在训练脚本里加一个命令行参数来控制是否打印这些调试信息，方便在不同阶段切换。

4.4 模型保存与加载的几种姿势

模型保存看起来简单，但里面有不少门道。最推荐的方式是保存state_dict，也就是模型的参数字典，而不是整个模型对象。因为整个模型对象依赖于具体的类定义，如果以后你改了类名或者文件结构，加载就会失败。保存state_dict的代码是torch.save(model.state_dict(), 'model.pth')，加载的时候先实例化模型，再调用model.load_state_dict(torch.load('model.pth'))。如果你还需要恢复训练，那就要把优化器的状态、epoch数、学习率调度器的状态也一起保存。我一般会保存一个字典，包含epoch、model_state_dict、optimizer_state_dict、scheduler_state_dict和best_acc。加载的时候用torch.load读进来，然后分别恢复。还有一个细节是map_location参数。如果你在GPU上训练，在CPU上加载，或者反过来，就需要指定map_location。比如torch.load('model.pth', map_location='cpu')。如果不指定，可能会报错。另外，如果你用了多GPU训练，保存的state_dict里的参数名会带module.前缀，加载到单GPU模型时需要去掉这个前缀。我一般会在保存之前用model.module.state_dict()来避免这个问题，或者加载的时候用OrderedDict做键名替换。

5. 常见问题与排查技巧实录

5.1 显存不足的排查与优化策略

显存不足是训练深度学习模型时最常见的问题之一。报错信息通常是CUDA out of memory。排查思路是这样的：首先看batch size是不是太大了，这是最直接的原因。你可以试着把batch size减半，如果问题解决，那就是显存不够。但减小batch size可能会影响训练效果，所以更好的办法是优化显存使用。第一个优化点是使用混合精度训练。pytorch内置了torch.cuda.amp，可以自动把部分计算转换成float16，显存占用能减少将近一半，而且速度还有提升。用法很简单，用autocast上下文管理器包裹前向传播，用GradScaler来缩放损失。第二个优化点是及时释放不需要的张量。比如在计算完loss之后，如果某个中间变量不再需要，可以用del删除，然后调用torch.cuda.empty_cache()。但要注意，频繁调用empty_cache会拖慢速度，所以只在关键时刻用。第三个优化点是检查是否有张量在计算图里累积。比如你在循环里不断把loss累加到一个列表里，却没有detach，那整个计算图都会保留在显存里。正确的做法是累加loss.item()或者loss.detach()。第四个优化点是使用梯度累积。如果你想要大的batch size但显存不够，可以分多次前向传播，累积梯度，然后再更新参数。这样等效于大batch训练，但显存占用不变。

5.2 损失不下降或出现NaN的排查路径

损失不下降或者变成NaN，是另一个让人头疼的问题。我一般按照以下顺序排查。第一步，检查数据。有没有标签错误？有没有输入数据包含NaN或者Inf？我习惯在数据加载之后加一行assert not torch.isnan(inputs).any()，虽然会稍微慢一点，但能提前发现问题。第二步，检查学习率。学习率太大会导致loss震荡甚至发散，太小会导致loss下降极慢。我一般会先用一个很小的学习率跑几个batch，确认loss在下降，然后再逐步调大。第三步，检查损失函数。比如分类任务用了nn.CrossEntropyLoss，但模型输出又加了softmax，这就重复了，会导致梯度异常。第四步，检查梯度。用torch.nn.utils.clip_grad_norm_裁剪梯度，如果裁剪前的梯度范数非常大，说明可能有梯度爆炸。第五步，检查网络结构。有没有层忘了加激活函数？有没有维度不匹配导致广播错误？我遇到过一次，两个张量维度不匹配，但因为广播机制没有报错，结果计算出来的loss完全不对。后来养成了习惯，在关键步骤打印张量形状。第六步，如果用了预训练模型，检查输入预处理是否和预训练时一致。比如ImageNet预训练的模型要求输入归一化到特定均值和方差，如果你忘了做，loss可能不下降。

5.3 训练速度慢的常见原因与加速手段

训练速度慢可能由多种原因造成。首先用torch.cuda.synchronize()配合time.time()来定位瓶颈在数据加载、前向传播还是反向传播。如果数据加载是瓶颈，可以增加num_workers，开启pin_memory，或者把数据预处理移到GPU上做。如果前向传播慢，检查模型里有没有低效的操作，比如用循环代替矩阵运算，或者频繁在CPU和GPU之间拷贝数据。如果反向传播慢，可能是模型太深或者batch size太大。加速手段方面，混合精度训练是最有效的，通常能提速30%到50%。另外，使用torch.backends.cudnn.benchmark = True可以让cudnn自动寻找最优的卷积算法，但前提是你的输入尺寸是固定的。如果输入尺寸变化，这个设置反而会拖慢速度。还有，如果模型里有大量的全连接层，可以考虑用torch.nn.Linear的bias=False来减少参数。对于卷积层，使用groups参数做分组卷积也能减少计算量。我实测过，在一个图像分类模型上，把标准卷积换成深度可分离卷积，速度提升了将近一倍，精度只掉了不到一个百分点。最后，如果单卡不够用，可以考虑用nn.DataParallel或者nn.parallel.DistributedDataParallel做多卡训练。DataParallel用起来简单，但效率不如DistributedDataParallel，后者需要每个进程独立启动，但扩展性更好。

5.4 常见报错速查表

提示：遇到报错不要慌，先看报错信息的最后一行，那里通常指出了具体的文件和行号。然后打印相关变量的形状和类型，大部分问题都能定位。

6. 从研究到部署的过渡经验

6.1 TorchScript转换的实操要点

当你研究阶段结束，准备把模型部署到生产环境时，TorchScript是一个绕不开的工具。它可以把你的动态图模型转换成静态图表示，然后序列化保存，脱离Python环境运行。转换方式有两种：torch.jit.trace和torch.jit.script。trace的方式是喂一个示例输入，记录下所有执行过的操作。这种方式简单，但有个致命缺点：它只能记录实际执行到的路径。如果你的模型里有if-else分支，trace只会记录示例输入走的那条分支，其他分支会被丢弃。所以对于有控制流的模型，必须用script方式。script方式会直接分析你的Python代码，把它编译成TorchScript。但script对Python语法的支持有限，比如不支持某些动态特性，不支持*args和**kwargs的某些用法。我一般会先尝试script，如果报错太多，再退回trace，同时确保模型里没有依赖输入值的控制流。转换完成之后，一定要用torch.jit.save保存，然后在C++或者Python的无梯度环境下加载测试，确保输出和原模型一致。我踩过的坑是：转换后的模型在CPU上运行正常，但在GPU上因为某些算子不支持而报错。所以转换后要在目标设备上完整测试一遍。

6.2 ONNX导出的注意事项

ONNX是另一个常用的模型交换格式，它可以让你的pytorch模型跑在TensorRT、OpenVINO等其他推理引擎上。导出ONNX用torch.onnx.export函数，需要提供模型、示例输入、导出路径、输入输出名称等参数。这里有几个关键点。第一，opset_version的选择。不同的推理引擎支持的opset版本不同，我一般会选一个较新的稳定版本，比如11或者12，然后根据目标引擎的文档调整。第二，动态维度的处理。如果你的模型需要支持变长输入，比如变长序列或者不同尺寸的图像，需要在导出时用dynamic_axes参数指定哪些维度是动态的。第三，导出后的验证。导出完成之后，一定要用onnxruntime加载模型，用相同的输入跑一遍，对比输出是否和pytorch一致。我遇到过导出成功但推理结果完全不对的情况，后来发现是某个自定义算子没有被正确转换。第四，注意模型里的操作是否都被ONNX支持。比如某些复杂的索引操作、自定义的autograd函数，可能无法导出。这时候就需要改写模型，用ONNX支持的操作来替代。我一般会在研究阶段就尽量使用标准算子，避免后期转换时大改。

6.3 部署时的性能优化技巧

模型部署到生产环境后，性能是首要考虑的因素。除了前面提到的混合精度和TorchScript，还有几个技巧。第一，算子融合。推理引擎通常会自动做算子融合，比如把卷积、批归一化和激活函数融合成一个算子，减少内存访问和计算开销。你不需要手动做，但要知道这个机制的存在。第二，量化。把float32的权重和激活值转换成int8，可以大幅减少模型大小和推理延迟。pytorch提供了torch.quantization模块，支持动态量化和静态量化。动态量化最简单，只需要一行代码，适合LSTM和全连接层。静态量化需要校准数据，精度损失更小，适合卷积层。我实测过，在一个图像分类模型上做静态量化，模型大小缩小到原来的四分之一，推理速度提升了一倍多，精度只掉了不到一个百分点。第三，批处理。在生产环境中，尽量把多个请求合并成一个batch一起推理，这样可以充分利用GPU的并行能力。但要注意延迟和吞吐量的权衡，batch太大会导致单个请求的响应时间变长。第四，使用TensorRT。如果你的部署环境是NVIDIA GPU，TensorRT可以进一步优化推理性能。它会对模型进行层融合、精度校准、内核自动调优等操作。从pytorch导出ONNX再转TensorRT是常见的路径，但转换过程中可能会遇到不支持的算子，需要自己写插件。

7. 一些零散但重要的经验补充

7.1 随机种子与可复现性

做实验的时候，可复现性非常重要。pytorch的随机性来源有好几个：权重初始化、数据打乱、Dropout、数据增强。要固定随机种子，需要设置torch.manual_seed(seed)、torch.cuda.manual_seed_all(seed)、np.random.seed(seed)和random.seed(seed)。另外，还要设置torch.backends.cudnn.deterministic = True和torch.backends.cudnn.benchmark = False。但要注意，这样设置会牺牲一些性能，因为cudnn不再自动寻找最优算法。我一般会在实验阶段开启确定性，在最终跑生产模型时关闭。还有一个细节是，即使设置了所有种子，不同版本的pytorch或者不同的GPU型号，结果也可能有微小差异。所以可复现性是相对的，不要期望完全一致。我一般会跑多次实验取平均值，而不是依赖单次结果。

7.2 学习率查找器的使用

学习率是最重要的超参数之一。手动调学习率很费时间，我推荐用学习率查找器（LR Finder）来快速找到一个合适的范围。原理很简单：从一个很小的学习率开始，每个batch指数级增加学习率，同时记录loss。然后画出loss随学习率变化的曲线，选择loss下降最快的那段学习率。pytorch没有内置LR Finder，但可以自己实现，也可以用现成的库。我一般会选曲线最低点对应的学习率，或者最低点左边一点的值。这个方法能帮你省去大量盲目尝试的时间。找到合适的学习率之后，再配合学习率调度器，比如CosineAnnealing或者OneCycleLR，往往能取得更好的效果。OneCycleLR是我最近比较喜欢用的，它让学习率先上升再下降，训练初期快速收敛，后期精细调整。实测下来，用OneCycleLR比固定学习率收敛快不少，最终精度也更高。

7.3 梯度累积与大批量训练的等效实现

前面提到过梯度累积，这里展开说一下具体实现。假设你想用batch size 128训练，但显存只够放32。你可以把128个样本分成4个mini-batch，每个32，依次前向传播计算loss，然后loss除以4再反向传播，这样梯度就累积了4次。等4个mini-batch都处理完，再调用optimizer.step()更新参数，然后optimizer.zero_grad()清空梯度。这样等效于batch size 128的训练。代码上就是在循环里加一个判断，每累积N次才更新一次。要注意的是，如果你用了BatchNorm，梯度累积会影响统计量的计算，因为BatchNorm是在每个mini-batch上计算均值和方差的。所以梯度累积和BatchNorm一起用的时候要小心，可能需要用同步BatchNorm或者改用GroupNorm。我一般在梯度累积的时候会把BatchNorm的momentum调小一点，让统计量更新更平滑。另外，学习率也需要相应调整，因为等效batch size变大了，学习率通常也要按比例增大。

7.4 自定义autograd函数的场景与写法

大多数情况下你不需要自定义autograd函数，因为pytorch已经覆盖了绝大多数操作。但有些特殊场景，比如你想实现一个非标准的前向和反向计算，或者想用numpy实现前向但需要pytorch的反向，就需要自定义torch.autograd.Function。写法是继承torch.autograd.Function，实现forward和backward两个静态方法。forward里做计算，backward里根据上游梯度计算对输入的梯度。关键点是forward里如果需要保存中间结果用于反向，要用ctx.save_for_backward保存，然后在backward里用ctx.saved_tensors取出来。我遇到过一个坑：在forward里返回的张量如果是从numpy数组转换来的，需要确保它是torch.Tensor类型，并且requires_grad设置正确。另外，backward的返回值个数必须和forward的输入个数一致，不需要梯度的输入返回None。自定义autograd函数一般用于实现一些特殊的损失函数或者复杂的数学运算，普通模型开发很少用到。但了解它的机制，能帮你更好地理解pytorch的自动求导原理。

7.5 分布式训练入门：DataParallel与DistributedDataParallel

当单卡显存不够或者训练太慢时，就需要多卡训练。pytorch提供了两种方式：nn.DataParallel和nn.parallel.DistributedDataParallel。DataParallel最简单，只需要把模型用nn.DataParallel(model)包一层，然后正常训练就行。但它有个缺点：每个batch的数据会被切分到各张卡上，模型会在每张卡上复制一份，梯度汇总到主卡上更新，然后再广播到其他卡。这个过程主卡负载高，效率不高。DistributedDataParallel是更推荐的方式，它每个进程独立控制一张卡，梯度通过NCCL后端做all-reduce，效率更高。但配置起来麻烦一些，需要设置init_process_group，用DistributedSampler来切分数据，启动方式也从python train.py变成torchrun --nproc_per_node=4 train.py。我一般会在单卡上把模型调通，然后再迁移到DistributedDataParallel。迁移过程中要注意：保存和加载state_dict时，参数名会带module.前缀，需要处理；日志只在主进程打印，避免重复输出；BatchNorm的统计量需要同步，可以用torch.nn.SyncBatchNorm来替换普通的BatchNorm。

8. 我个人在实际操作中的体会

写了这么多，其实最想说的是：pytorch这个工具，入门容易，精通难。它的API设计很直观，你花一个下午就能跑通一个MNIST分类。但真正要把模型训好、训稳、部署好，需要踩的坑一个都不会少。我这些年最大的体会是，不要迷信任何“最佳实践”，包括我上面写的这些。每个项目的数据、任务、硬件环境都不一样，你需要根据实际情况去调整。比如学习率，别人用0.001效果好，你的数据可能用0.0001才收敛。比如优化器，Adam在大多数情况下够用，但有些任务SGD就是比Adam泛化好。我的建议是，先把baseline跑通，然后每次只改一个变量，观察效果。记录每次实验的配置和结果，慢慢你就有了自己的直觉。另外，多读pytorch的官方文档和源码，很多问题文档里其实都写了，只是你没注意到。最后再分享一个小技巧：如果你在训练过程中发现loss突然变成NaN，不要急着从头开始，先检查是不是某个batch的数据有问题。我遇到过一次，数据集中有一张图片损坏了，读出来全是NaN，导致整个训练崩溃。后来在__getitem__里加了异常捕获，跳过损坏的样本，问题就解决了。这种小坑，只有真正跑过项目的人才会懂。