ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

模型优化实战指南:从优化器选型到TensorRT加速

模型优化实战指南:从优化器选型到TensorRT加速 上个月我调一个图像分类模型ResNet-50 的骨架中等量级的数据集。起手式很常规PyTorch 默认的 Adamlr 设 1e-3。结果训练日志特别难看——前 20 个 epoch 还像样到了 40 epoch验证集 loss 卡在 0.5 附近再也不动。我把数据增强、模型结构全都检查了一遍最后抱着试试看的心态把优化器换回 SGD momentum同一个 batch size、同一个学习率调度loss 直接压到 0.32。这件事让我意识到一件事很多人口中的 Model-Optimizer 并不是某个具体工具的名字而是覆盖训练与推理两侧的一整套模型优化链路——训练时怎么选 optimizer、怎么配参数推理前怎么压缩模型、怎么加速算子。这篇文章就把我这套完整落地流程、以及踩过的坑原原本本写出来。适合正准备把模型从能跑折腾到跑得快、跑得稳的工程师参考也适合刚入行、对优化器选型和模型加速没什么概念的同学当一份作业抄。1. 训练端的优化器选型为什么 Adam 不是灵丹妙药1.1 从一次翻车说起默认参数适配一切是个错觉先说那次翻车。模型是 ResNet-50优化器用的 torch.optim.Adam参数完全默认lr1e-3betas(0.9, 0.999)weight_decay0。40 个 epoch 下来 loss 卡在一个平坦区域训练集上 loss 也不怎么降。那段时间我一度怀疑是数据问题换了更强增强、加了 mixup都没明显起色。后来我重新翻了几篇经典论文发现自己踩了一个很常见的坑Adam 的默认参数是为大规模稀疏梯度、Transformer 类任务调出来的不是为稠密 CNN 训练调出来的。CNN 的梯度相对稠密、但带有比较明显的方向性SGD momentum 在这些任务上往往能走出一条更稳定的泛化曲线。Adam 的全局自适应学习率会让参数在一个平坦区域附近反复震荡尤其是 lr1e-3 对 CNN 来说通常偏高很容易越过最优谷。换回 SGD 之后的参数我只改了三个momentum0.9、base_lr0.1配合 cosine 调度、weight_decay1e-4。同样算力、同样 epoch 数验证集 loss 从 0.5 压到 0.32top-1 精度也涨了约 1.7 个点。这次之后我给自己定了个规矩不经过验证的默认 optimizer 只配跑通 demo不配跑正式实验。1.2 按任务特征选优化器SGD、Adam、AdamW、LAMB 到底怎么选选优化器不能只看哪个最新用哪个要看任务特征。我平时主要看三个维度模型结构、数据规模、训练成本。优化器适用场景常用参数我的使用频率SGD MomentumCNN、中小规模数据、追求泛化精度lr0.1~0.3momentum0.9wd1e-4CV 任务跑榜首选Adam小数据快速验证、RNN/早期实验lr1e-3 起步需配合调度只用来做 ablationAdamWTransformer/BERT/LLM、decoupled weight decaylr3e-5~1e-3wd0.01~0.05主力训练语言模型和 ViT 都靠它LAMB超大 batch上千预训练、BERT/大规模对比学习与 AdamW 类似增加逐层自适应大规模集群训练才用这里有个具体的判断逻辑如果任务是 CNN 图像分类、目标检测这类稠密结构SGD momentum 依然是性价比最高的选择尤其是 ResNet、RegNet、RepVGG 这类结构SGD 的泛化性明显好于 Adam。如果任务是 Transformer 类AdamW 基本是必须的——它把 weight decay 从梯度里解耦出来能独立控制参数衰减这是 Adam 做不到的。如果数据量很小、只想快速看模型能不能收敛我会先用 Adam 以 1e-3 的 lr 跑 10 个 epoch观察 loss 下降曲线是否稳定再用正式优化器拉长训练。LAMB 我目前只在预训练阶段用它把每个 layer 的学习率做了自适应缩放能让上千的 batch size 稳定收敛但配套的 warmup 策略也得更精细否则前期特别容易炸。1.3 warmup、weight decay、betas优化器的三组配套参数选好优化器只是第一步参数不配套等于白选。我总结下来真正影响收敛质量的通常是三组参数。第一组是 warmup。Transformer 和超大 batch 训练必须加 warmup原因很简单训练初期参数离最优解非常远梯度方向噪声极大如果一开始就用大学习率容易冲进一个坏的损失谷后面再也出不来。我一般的做法是 linear warmup占整个训练步数的 5%~10%。对大规模预训练任务这个比例可以提到 10% 以上。第二组是 weight decay。这里有个特别容易混淆的点Adam 里的 weight_decay 实际是 L2 正则项等于把衰减项加到了梯度里而 AdamW 把它变成了直接对参数做衰减二者数学上不等价。这就是为什么很多人把 SGD 的 weight decay 习惯直接搬进 AdamW 后效果很差。我通常的做法Transformer 用 AdamW 时 wd0.01CNN 用 SGD 时 wd1e-4CLIP 类对比学习会稍微加大到 0.05。第三组是 betas。beta1 控制一阶动量0.9 是几十年的老经验基本不动beta2 控制二阶动量默认 0.999 适合长程平稳训练但如果你的梯度噪声比较大比如 GAN、对比学习、小 batch 训练我会把 beta2 下调到 0.95~0.98让自适应学习率对短期波动更敏感收敛也会更稳。一次实际配置的代码长这样optimizer torch.optim.AdamW( model.parameters(), lr3e-4, betas(0.9, 0.999), eps1e-8, weight_decay0.01, ) total_steps len(train_loader) * epochs scheduler get_cosine_schedule_with_warmup( optimizer, num_warmup_stepsint(total_steps * 0.05), num_training_stepstotal_steps, )提示transformers 库里的 AdamW 和 torch.optim.AdamW 严格来说实现细节不同前者的 weight decay 是解耦式的、对学习率不敏感后者接近原论文。换库时最好重新验证一次收敛曲线不要直接默认等价。2. 推理端的压缩顺序先体检、再减脂、后量化2.1 动手之前先给模型做一次体检很多同学拿到模型就想立刻上 TensorRT、立刻量化结果性能没提升多少精度掉了不少。我不一样我习惯先花半小时给模型做体检确认瓶颈到底在哪。体检工具用 torch.profiler 就够了能拿到 CPU/GPU 时间占比、算子耗时、显存占用from torch.profiler import profile, ProfilerActivity with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapesTrue, ) as prof: model(input_tensor) print(prof.key_averages().table(sort_bycuda_time_total, row_limit20))我一般盯四个指标参数量、FLOPs、单次推理延迟、显存占用。参数量大但延迟不高说明算子并行度好参数量中等但延迟高往往是算子碎片化、kernel launch 开销过大如果 FLOPs 高得离谱就考虑结构化剪枝如果激活值缓存吃显存就要往量化和算子融合方向上想。有一次我拿到一个检测模型在 A10 上跑单张图要 25msprofiler 一查发现 60% 时间花在若干个很小的 reshape 和 transpose 上。这种算子本身 FLOPs 极低但每次都要做 kernel launch慢就慢在调度上。这种情况下量化收益有限真正该做的是把模型重写、把多个小算子合并成一个大算子或者直接让编译器来融合。2.2 量化精度损失的常见根源与应对量化是目前最有效的推理加速手段之一把 FP32 模型压缩成 FP16、INT8计算带宽和显存占用都能成倍下降。但量化掉精度也是常态尤其 PTQ训练后量化最容易出问题。我踩过的第一个坑是 BN 层没有融合。PyTorch 里 BN 在 eval 状态下会采用 running stats而不是按当前 batch 重新统计量化仿真时如果保留独立的 BN 层数值分布会产生偏移。正确做法是在导出和量化之前把 BN 层 fold 进前面的 Conv 层PyTorch 可以通过 torch.fx 配合 fusing 工具实现或者先冻结 BN、再用 QAT。第二个坑是校准数据集数量太少。PTQ 会收集激活值的统计分布来截断量化阈值如果只拿几十张图校准阈值很容易被偶然出现的异常值带偏。我的经验是最少拿 500~1000 张具有代表性的图类别分布要接近真实上线场景。数据量不够时宁可多跑几分钟校准也不要省这一步。第三个坑是 per-tensor 和 per-channel 的选择。权重量化强烈建议 per-channel因为分布差异在不同 channel 之间可能很大per-tensor 会让某些 channel 的精度严重受损激活值因为要支持矩阵计算用 per-tensor 更通用。对称量化在激活值大部分为正、比如 ReLU 输出场景下效果好但如果激活值分布有负值、且偏移比较明显非对称量化更稳。如果 PTQ 精度损失超过 1%不要反复调阈值硬撑直接换 QAT。QAT 相当于在训练中插入 fake 量化节点让模型自己适应量化误差。代价是多跑一轮微调但结果通常比 PTQ 稳定得多。2.3 知识蒸馏的调配技巧与和剪枝的先后顺序蒸馏是精度补偿的重要手段。一个训好的大模型当 teacher去带小的 student学生模型模仿教师输出的 soft logits通常比只用 one-hot 标签收敛得更快更准。我用的标准蒸馏 loss 是import torch.nn.functional as F def distill_loss( student_logits, teacher_logits, hard_labels, temperature5.0, alpha0.7, ): kl_loss F.kl_div( F.log_softmax(student_logits / temperature, dim-1), F.softmax(teacher_logits / temperature, dim-1), reductionbatchmean, ) * (temperature ** 2) ce_loss F.cross_entropy(student_logits, hard_labels) return alpha * kl_loss (1 - alpha) * ce_loss温度 T 我一般取 3~7。T 太小soft label 几乎退化成 one-hot蒸馏没意义T 太大teacher 输出的概率分布被抹平变成噪声。alpha 控制在 0.7 左右比较安全给 hard label 保留一定权重防止 student 完全复读 teacher 的错误判断。至于剪枝和蒸馏谁先谁后我现在的固定顺序是先做结构化剪枝确定通道数再用 teacher 蒸馏补偿精度最后量化部署。先剪枝的好处是 student 的结构被提前锁定蒸馏出来的模型可以直接用于后续量化不会出现蒸馏完发现通道太宽、剪完又要重新蒸馏的返工。而如果 teacher 模型大到显存放不下那就只能先蒸馏压缩出一个小模型再对这个模型做轻量级剪枝最后微调。提示结构化剪枝尽量选择 channel pruning 而不是稀疏化。稀疏化虽然能压参数量但推理硬件很难吃到收益GPU、TensorRT 对稀疏矩阵的支持都很有限。3. 从训练到部署的完整链路PyTorch、ONNX 与 TensorRT 的取舍3.1 训练阶段的配套细节梯度裁剪、EMA 与混合精度优化器选好只是热身真正让训练收敛质量上台阶的是几个配套动作。梯度裁剪是我开箱必用的。Transformer 训练、以及所有用了大学习率任务的训练都可能在某个 batch 里遇到异常梯度不裁剪的话 loss 直接飞向 NaN。我统一用全局范数裁剪阈值 1.0这个值适配大部分模型。常用写法torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)EMA 指数移动平均是提精度的隐藏神器。不直接训练模型参数而是训练过程中持续维护一组参数影子用来做评估和部署。虽然多占一份显存但效果几乎总是正面的尤其对检测、分割这类任务EMA checkpoint 比普通 checkpoint 稳定不少。我常用的衰减率是 0.999如果你训练步数短、数据噪声大可以调到 0.99。混合精度训练我基本无脑开。autocast GradScaler 的组合能省接近一半显存训练速度也有明显提升。要注意的是一旦开了 FP16loss scaling 的初始值、是否触发 scale backoff 都要看日志别让 underflow 悄悄毁掉精度scaler torch.cuda.amp.GradScaler() for step, (x, y) in enumerate(loader): optimizer.zero_grad() with torch.cuda.amp.autocast(): loss criterion(model(x), y) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update()3.2 导出 ONNX 最容易卡住的四个问题训练完的 PyTorch 模型不能直接部署到生产环境第一步通常是导出 ONNX。这个过程我踩过不少坑按出现频率排序第一忘记 model.eval()。导出时模型处于 train 模式Dropout 和 BN 的随机性全被冻结进图里推理结果抖动。这个排查起来很容易但发生频率最高。第二opset 版本不够。某些新算子像 adaptive_avg_pool、grid_sample、一些 attention 实现老版本 opset 根本不认识。我现在的默认值是最低 14推荐 17。版本越高表达能力越强但也可能因为太新导致推理引擎不支持要两头兼顾。第三dynamic shape 配置错误。固定输入尺寸的模型导出最容易但一旦上线需要不同分辨率就要给 dynamic axes 配置合适的范围。TensorRT 优化引擎时也要额外提供 minShape、optShape、maxShape配置不对会出现引擎构建成功但每次推理都 rebuild的情况速度反而更慢。第四算子不支持。某个自定义 op 在 ONNX 里找不到对应实现要么升级 opset要么把自定义 op 用基础算子重写。我建议导出前就查一下算子映射表别等导完跑起来才报错。导出完成后我有个强制习惯先用 onnxruntime 跑一遍和 PyTorch 输出做数值对齐检查import onnxruntime as ort import numpy as np sess ort.InferenceSession(model.onnx, providers[CPUExecutionProvider]) onnx_out sess.run(None, {input: x_np})[0] torch_out model(x_torch).detach().cpu().numpy() max_diff np.max(np.abs(torch_out - onnx_out)) print(fmax diff: {max_diff:.2e})差值超过 1e-4 就必须停下来查算子精度问题不要带着误差继续往 TensorRT 走。3.3 TensorRT 与 torch.compile部署场景怎么选推理引擎这块我同时维护两套方案因为使用场景完全不同。TensorRT 是 NVIDIA GPU 上的生产首选。它会把 GPU kernel 做算子融合、内存复用、精度校准FP16 下常见的 100ms 级模型能压到 30~40msINT8 还能再降一半。代价是构建时间长、引擎文件与 CUDA/驱动版本强绑定、动态 shape 配置繁琐。适合对延迟要求高的 API 服务。torch.compile 是 PyTorch 2.x 的原生方案。开发迭代期特别爽不用导出模型、不用换 runtime一行代码就能让模型提速 30%~50%。但它本质还是依赖 PyTorch 运行时不适合做对外交付的轻量产物。我的用法是实验验证阶段用 torch.compile 确认提速空间有多大正式部署再换 TensorRT。ONNX Runtime 更适合跨平台交付CPU/GPU 都支持还能挂 TensorRT execution provider等于多一层兼容。如果你的客户环境有各种 GPU、CPU 混用ONNX Runtime TRT EP 是最稳的组合。我自己的选择逻辑很简单单模型高频推理用 TensorRT FP16快速原型直接 torch.compile需要跨环境交付的用 ONNX Runtime。三者不是替代关系是一条流水线上的不同工位。4. 一次完整的 Model-Optimizer 落地过程实测数据与典型坑4.1 一组优化前后的实测对比拿我之前做过的一个 ResNet-50 分类模型举例数据是 ImageNet 子集精度数值是我的实际记录不代表所有环境阶段Top-1 精度延迟(ms)模型体积备注PyTorch FP32 baseline77.212.598MB默认 SGD 训练换 AdamW EMA cosine77.812.498MB训练端优化ONNX Runtime FP3277.89.898MB去掉框架开销TensorRT FP1677.64.249MB算子融合半精度TensorRT INT8 (PTQ)76.82.725MB校准集 1000 张每一步的收益逻辑是训练端调整白拿了约 0.6 个点ONNX 导出省掉 PyTorch 框架调度开销TensorRT 的算子融合、FP16 半精度砍掉一半带宽INT8 再把延迟和体积压到接近 1/5。但精度和延迟的 trade-off 非常明显跌到 77% 以内的收益我就不想要了所以最终线上用 FP16 版本。4.2 先找瓶颈再动手别把时间花在无效优化上有些模型慢在数据加载和预处理有些模型慢在小算子调度如果一上来就套量化很容易白忙一场。判断瓶颈有一个土办法先用 nvidia-smi 看 GPU 利用率。如果训练或推理时 GPU util 长期低于 50%先别碰模型去查数据迭代、预处理 pipeline、以及 CPU 和 GPU 之间的拷贝开销。我遇到过推理服务端每次请求都重新加载 engine、做完推理就释放显存上下文的情况GPU util 只有 20%但模型的推理算子本身并不慢。改成常驻 context 后延迟从 28ms 降到 2.5ms十倍差距根本不是模型算不动是服务端实现烂。如果 GPU util 很高但延迟还是压不下去再用 profiler 拉算子耗时表找到 Top 3 的耗时算子针对性地做融合、重写或量化。有一类问题很典型模型里有一堆零散的 slice、transpose每个耗时不足 0.1ms但数量上百个叠加起来就成了主要瓶颈。这种场景直接让 TensorRT 做图优化往往比手动改 PyTorch 代码高效得多。4.3 三个真实踩坑记录每一个都值一次返工第一个坑量化后精度暴跌。当时我在做一个小目标检测模型PTQ 校准后精度掉了 4 个点非常反常。最后定位到原因模型里 BN 层没有 fold 到 Conv 里量化仿真把 BN 独立计算时的激活分布统计错了。解决方法是先把 BN 层融合进 Conv再重新校准精度回到 97% 以上。这个坑在 ResNet、EfficientNet 这类带 BN 的模型上特别容易出现。第二个坑蒸馏温度调太高。有次做一个小 Transformer 蒸馏我把温度设成 10alpha 设成 0.9结果 student 在验证集上比不蒸馏还差。原因是温度太高后 teacher 的 softmax 输出几乎接近均匀分布软标签变成噪声alpha 太高又把 hard label 的真实信号盖住了。后来改成 T5、alpha0.7精度一下子回来了。教训是蒸馏不是参数越极端越好要对齐 teacher 置信度水平。第三个坑TensorRT 动态 shape 配置太宽。我为了上线不同的输入分辨率把 height 和 width 的动态范围设成了 1~4096结果引擎构建时间超长而且小尺寸输入的实际性能比固定 batch 差很多。TensorRT 的 auto-tuning 是按 shape 做的范围太宽每个 shape 都选不到最优 kernel。后来我把范围压缩到实际业务区间再提供 minShape、optShape、maxShape 三档速度和构建时间都恢复正常。现在我做模型优化已经固定成一套流程先打基线再做训练端调整然后推理端压缩最后按环境选引擎。每一步只改一个变量并且留一份对照日志。这个习惯帮我省掉了至少三次返工也让我在向同事解释为什么这一步要这么做的时候永远有数据可讲。你如果也在折腾模型优化建议也先建一个这样的基线表格后续每一步的收益和损失都会很直观。
返回列表