资讯详情

资讯详情

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_gradTrue然后所有基于这个张量的运算都会被记录最后调用.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_norm1.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(weightclass_weights)来给少数类更高的权重。优化器的选择上Adam是我最常用的默认选项因为它对学习率不那么敏感收敛也快。但如果你追求极致的泛化性能SGD加上动量往往能取得更好的结果只是需要更精细的学习率调度。我一般会先用Adam快速跑一个baseline确认模型结构没问题然后再换SGD调参。学习率调度器也很重要StepLR、CosineAnnealingLR、ReduceLROnPlateau我都用过。ReduceLROnPlateau比较省心它根据验证集loss自动降低学习率不需要你手动设定下降的时机。但要注意它的patience参数设得太小会导致学习率降得太快模型还没收敛就停了。4.3 训练循环的骨架与关键监控指标一个标准的训练循环包含以下几个步骤前向传播、计算损失、反向传播、更新参数。我习惯把训练和验证写成一个函数接受模型、数据加载器、优化器、损失函数等参数。在训练循环里我会记录每个epoch的平均loss和准确率。准确率的计算要注意如果是多分类需要用torch.argmax(outputs, dim1)拿到预测类别然后和标签比较。这里有个细节在计算准确率的时候记得用.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_locationcpu)。如果不指定可能会报错。另外如果你用了多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的biasFalse来减少参数。对于卷积层使用groups参数做分组卷积也能减少计算量。我实测过在一个图像分类模型上把标准卷积换成深度可分离卷积速度提升了将近一倍精度只掉了不到一个百分点。最后如果单卡不够用可以考虑用nn.DataParallel或者nn.parallel.DistributedDataParallel做多卡训练。DataParallel用起来简单但效率不如DistributedDataParallel后者需要每个进程独立启动但扩展性更好。5.4 常见报错速查表报错信息可能原因解决方法RuntimeError: Expected all tensors to be on the same device张量或模型不在同一设备上检查所有张量和模型是否都.to(device)RuntimeError: size mismatch张量维度不匹配打印相关张量的.shape检查全连接层输入维度RuntimeError: element 0 of tensors does not require grad损失函数输入没有梯度检查模型参数是否requires_gradTrue是否用了detach()CUDA out of memory显存不足减小batch size使用混合精度及时释放张量ValueError: Expected input batch_size to match target batch_size输入和标签的batch维度不一致检查DataLoader的batch_size和drop_last设置RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation原地操作破坏了计算图避免使用带下划线的原地操作或改用非原地版本AttributeError: DataParallel object has no attribute xxx多卡训练时访问了自定义属性用model.module.xxx访问或把属性定义在forward里提示遇到报错不要慌先看报错信息的最后一行那里通常指出了具体的文件和行号。然后打印相关变量的形状和类型大部分问题都能定位。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 GPUTensorRT可以进一步优化推理性能。它会对模型进行层融合、精度校准、内核自动调优等操作。从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_node4 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__里加了异常捕获跳过损坏的样本问题就解决了。这种小坑只有真正跑过项目的人才会懂。
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →