简介深度学习模型的训练与部署之间存在一条需要精细打通的链路。模型在PyTorch中训练收敛后生产环境往往要求更高的吞吐与更低的延迟。ONNX作为开放的模型交换格式能够将PyTorch计算图无缝转移至NVIDIA TensorRT引擎借助层融合、FP16/INT8精度校准与内存复用等优化手段在GPU上实现数倍推理提速。以经典的mobilenetv2图像分类模型为例其深度可分离卷积结构在Jetson和服务器显卡上均有良好适配性是验证部署全流程的理想选择。实践中需要关注版本兼容性、预处理对齐、动态batch设置以及显存缓冲区复用等工程细节避免精度掉点和性能瓶颈。掌握这套方法论后可轻松迁移到YOLO等更复杂模型让深度学习应用真正高效落地。 先说个场景模型在 PyTorch 里训练得好好的准确率 95% 往上结果一上生产环境CPU 推理一帧几百毫秒GPU 上用原生 PyTorch 跑也就勉强几十毫秒吞吐完全扛不住。这时候你大概率会想到 TensorRT。我这次就用 mobilenetv2 走了一遍“PyTorch 训练 - ONNX 导出 - TensorRT 部署”的完整链路把过程中踩过的坑、验证过的细节、能直接抄作业的配置都整理出来。这篇文章适合刚接触部署、或者已经在用 PyTorch 但想把手头分类模型真正落地到 NVIDIA GPU 上的同学内容覆盖训练、导出、转换、推理四个环节重点是让你少走弯路。mobilenetv2 这个模型选型本身就很有意思它不算最新但足够经典深度可分离卷积结构练出来的模型体积小、速度快在 Jetson 和普通显卡上都有着很好的泛用性加上 PyTorch 官方 torchvision 里直接有预训练权重拿来做部署链路验证再合适不过。1. 整体链路设计与方案选型1.1 为什么选 mobilenetv2 TensorRT 这条路线先说部署方案选型的底层逻辑。做图像分类服务主流选择无非是几类直接用 PyTorch 的 TorchScript 做生产推理用 ONNX Runtime 跑 ONNX 模型或者用 TensorRT 把模型编译成高度优化的 engine。这三条路我都实际走过说下各自感受。TorchScript 最省事导出之后还是 PyTorch 运行时环境依赖几乎不用变但推理性能提升非常有限它解决的更多是“部署环境没有训练代码”的问题而不是“推理要更快”的问题。ONNX Runtime 是中间态兼容性好CPU 和 GPU 都能跑性能比原生 PyTorch 好一截但要真正榨干 GPU 算力TensorRT 依然是最优解——它会在构建 engine 时做层融合、精度校准、内存复用对 mobilenetv2 这种结构规整的 CNN 来说收益非常明显。我再具体举一个数字。同一张 224x224 输入在 RTX 3060 上用 PyTorch CUDA 跑 mobilenetv2 单张推理大约是 6 到 8 毫秒换成 TensorRT FP16 engine 之后能压到 2 到 3 毫秒连续 batch 推理场景下的吞吐差距更明显。对在线分类服务来说这个差距就是能否扛住峰值流量的区别。TensorRT 部署也确实有门槛比如版本匹配问题、动态 shape 的处理、精度校准流程但作为工程师这些属于一次投入、长期受益的能力。1.2 整体技术链路拆解整条链路可以分成四个清晰阶段我画不出图但用文字描述就是PyTorch 训练 - 导出 ONNX - TensorRT 构建 engine - 推理服务训练阶段保证模型精度导出阶段保证计算图正确转换阶段保证 TensorRT 能完整解析模型结构推理阶段保证预处理和后处理跟训练时完全对齐。每个阶段都有独立的验证方式也有各自的坑。先交代一下我的运行环境方便你对照组件版本操作系统Ubuntu 20.04 / Windows 11双环境测试GPUNVIDIA RTX 3060 12GB / Jetson Orin NanoPython3.8 / 3.10PyTorch1.13.0cu117 / 2.0.1cu118CUDA11.7 / 11.8cuDNN8.6.0TensorRT8.5.3.1Python APIONNX1.13.0onnxruntime-gpu1.14.0这里要提醒一句TensorRT 的版本和 CUDA、PyTorch 的版本耦合非常严重如果不想在环境配置上浪费一整天建议直接按官方文档的兼容矩阵对版本不要追求最新版稳定匹配才是第一位。2. PyTorch 训练阶段的关键细节2.1 环境搭建与训练脚本骨架训练阶段重点不是把 mobilenetv2 训到多高的精度而是要保证训练配置和后续部署推理完全一致。我这里用 PyTorch 官方的 torchvision 模型做迁移学习数据集是自定义的 10 类花分类数据。关于 PyTorch 环境搭建我多说一句GPU 版本安装前先查清楚自己的 CUDA 版本再选择对应的 PyTorch 安装命令nvidia-smi显示的 CUDA 版本和torch.version.cuda不一定一致前者是驱动支持的版本后者才是 PyTorch 实际使用的运行时版本两者之间兼容即可。数据集准备好后直接用torchvision.datasets.ImageFolder加载目录结构必须规范data/ ├── train/ │ ├── class1/ │ ├── class2/ │ └── ... ├── val/ │ ├── class1/ │ ├── class2/ │ └── ...训练脚本里最关键的一段是数据预处理必须和部署端的预处理完全一致否则精度在部署后会莫名掉几个点。mobilenetv2 在 ImageNet 上的标准预处理是transform_train transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) transform_val transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])很多人有个误区觉得部署时只要做ToTensor和Normalize就够了完全忽略 Resize 和 CenterCrop 的方式。我见过不少项目训练时用 256 resize 中心裁剪部署时直接拿原图 resize 成 224x224输入分布不一样精度自然对不上。2.2 mobilenetv2 模型加载与训练配置模型加载直接用 torchvision替换分类头即可import torchvision.models as models model models.mobilenet_v2(weightsmodels.MobileNet_V2_Weights.IMAGENET1K_V1) model.classifier[1] torch.nn.Linear(model.last_channel, num_classes)这里有个容易忽略的点mobilenet_v2的last_channel属性必须是 1280官方实现里最后一层是Conv2d(320, 1280, kernel_size1)很多人自定义模型时改了通道数却忘记同步修改last_channel轻则报维度错误重则模型加载不完整。如果你的分类数较少比如只有 10 类也可以用width_mult0.5或0.75进一步压缩模型但 TensorRT 加速比在标准 width 下表现最好建议先跑通全流程再考虑压缩。训练超参我是这么配的参数值优化器SGDmomentum0.9, weight_decay1e-4初始学习率0.01学习率策略CosineAnnealingLRT_max30Batch Size64Epoch30损失函数CrossEntropyLoss迁移学习场景下建议一开始就解冻全部层配合较低学习率微调而不是只训分类头。只训练分类头的话骨干网络的 feature 不一定适配你的数据分布最后精度很难上去。SGD 加 cosine 学习率在分类任务上依然是非常稳的组合Adam 适合快速试错但最终精度 SGD 通常还是略胜一筹。训练过程中保存模型时最好把model.state_dict()和完整的model都存一份方便后续导出。要记录验证集上最优的 top-1 准确率作为部署精度对比的基准线。3. 从 PyTorch 到 ONNX 的模型导出3.1 导出 ONNX 需要特别注意的几个参数PyTorch 转 ONNX 是整条链路中看似简单、实际最容易出问题的一步。核心代码如下import torch model.eval() dummy_input torch.randn(1, 3, 224, 224).cuda() torch.onnx.export( model, dummy_input, mobilenetv2.onnx, opset_version12, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size}, output: {0: batch_size} } )这里面有两个关键选择。第一是否要开 dynamic_axes。如果你的部署场景 batch size 固定或者只做 batch1 的在线推理强烈建议不要开动态 batch而是固定 batch 导出。固定 batch 的 engine 在 TensorRT 里可以省去很多维度推导的 overhead推理性能更好。如果你确实需要动态 batch那就必须把max_batch_size传给 TensorRT builder并且推理时每一帧都要做好 shape 对齐。我自己的做法是线上服务固定 batch1离线批量任务单独导出 batch32 的版本两个 engine 互不干扰。第二opset_version 的选择。TensorRT 8.x 对 ONNX opset 11 到 13 支持比较成熟opset 太高反而可能出现某些算子不兼容。我测试过 mobilenetv2 在 opset 12 下转换最干净opset 13 在某些 TensorRT 老版本上会多出一些不必要的子图拆分。如果你用的 TensorRT 版本比较新比如 8.6 以上用 opset 13 也没问题但建 engine 之前一定要用trtexec做一次快速验证避免部署环境里才发现算子不支持。3.2 导出后如何验证 ONNX 模型没有坏掉导出后的第一件事不是急着转 TensorRT而是验证 ONNX 模型的输出和 PyTorch 是否一致。这一步我用 onnxruntime-gpu 做一致性校验import numpy as np import onnxruntime as ort # 生成固定输入 x torch.randn(1, 3, 224, 224).cuda() with torch.no_grad(): torch_output model(x).cpu().numpy() x_cpu x.cpu().numpy() sess ort.InferenceSession(mobilenetv2.onnx, providers[CUDAExecutionProvider]) onnx_output sess.run(None, {input: x_cpu})[0] # 比较输出 diff np.abs(torch_output - onnx_output).max() print(fmax diff: {diff:.6f})正常情况下FP32 的 PyTorch 输出和 ONNX Runtime 输出最大误差应该在 1e-5 量级如果差异超过 1e-3说明导出过程有算子转换失败需要检查模型里是否有自定义层或特殊算子。mobilenetv2 全是标准卷积、BN、ReLU6、深度可分离卷积这种情况几乎不会发生一旦发生大概率是 PyTorch 版本和 ONNX 版本不兼容导致的算子映射问题。验证阶段还有个隐藏点PyTorch 模型的model.eval()模式必须设置否则 BN 层的 running_mean 和 running_var 不会冻结导出的 ONNX 计算图里 BN 数值都是错的推理精度会直接崩掉。实际上如果你忘了设 eval 模式导出的模型每个 batch 输出都可能不一样这个问题在很多项目里出现过。4. TensorRT 部署实战4.1 用 Python API 构建 TensorRT engineTensorRT 部署有两种常见方式直接用官方的trtexec命令行工具把 ONNX 转成 engine或者写 Python 脚本调用 TensorRT Python API。trtexec适合快速验证模型能不能转、性能大概什么水平但生产环境建议用 Python API因为可以精细控制 batch size、精度模式、工作空间大小。用 Python API 构建 engine 的核心逻辑import tensorrt as trt logger trt.Logger(trt.Logger.INFO) builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, logger) with open(mobilenetv2.onnx, rb) as f: assert parser.parse(f.read()) config builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30) config.set_flag(trt.BuilderFlag.FP16) # 构建 engine serialized_engine builder.build_serialized_network(network, config) with open(mobilenetv2_fp16.engine, wb) as f: f.write(serialized_engine)注意几个细节。.engine文件和 ONNX 不同它是和特定 GPU 架构、TensorRT 版本强绑定的。你在 RTX 3060 上构建的 engine 拿到 Jetson Orin 上大概率加载不了换 TensorRT 版本也会失效。所以生产环境一般在线下构建、线上加载或者干脆在容器启动时现场构建一次。set_memory_pool_limit控制的是 TensorRT 的工作空间上限。mobilenetv2 很小1GB 完全够用。设置太大会让 builder 在层融合时多做很多尝试编译时间变长收益可能微乎其微。FP16 开关如果只开 FP16 不做 INT8 量化构建过程中的精度损失通常非常小分类任务掉点一般控制在 0.5% 以内但推理速度能提升 30% 到 60%。如果你的业务对精度极其敏感可以先只跑 FP32确认 Baseline 准确后再切 FP16。4.2 前处理、推理与后处理的完整实现engine 构建好之后就是推理环节。TensorRT Python API 的推理流程和 PyTorch 差别很大核心是 bindings 缓冲区的管理。我整理了一个可以直接跑的推理脚本import numpy as np import tensorrt as trt import cv2 class TRTInference: def __init__(self, engine_path): self.logger trt.Logger(trt.Logger.WARNING) self.runtime trt.Runtime(self.logger) with open(engine_path, rb) as f: self.engine self.runtime.deserialize_cuda_engine(f.read()) self.context self.engine.create_execution_context() # 获取输入输出名称 self.input_name self.engine.get_tensor_name(0) self.output_name self.engine.get_tensor_name(1) # 分配缓冲区 self.input_shape (1, 3, 224, 224) self.output_shape (1, 10) self.input_dtype np.float32 self.output_dtype np.float32 def infer(self, image): # 前处理resize - crop - normalize - CHW img cv2.resize(image, (256, 256)) h, w img.shape[:2] startx (w - 224) // 2 starty (h - 224) // 2 img img[starty:starty224, startx:startx224] img img.astype(np.float32) / 255.0 mean np.array([0.485, 0.456, 0.406], dtypenp.float32) std np.array([0.229, 0.224, 0.225], dtypenp.float32) img (img - mean) / std img img.transpose(2, 0, 1) img np.ascontiguousarray(img)[np.newaxis, ...] # 输入输出缓冲区 input_buf np.zeros(self.input_shape, dtypeself.input_dtype) output_buf np.zeros(self.output_shape, dtypeself.output_dtype) # 传入 GPU input_buf[...] img # 这里为了简化展示用了 CPU 拷贝到 CUDA 内存。 # 实际生产建议用 pycuda 或 cuda-python 管理显存避免每帧分配 d_input cuda.mem_alloc(int(input_buf.nbytes)) d_output cuda.mem_alloc(int(output_buf.nbytes)) cuda.memcpy_htod(d_input, input_buf.tobytes()) self.context.set_tensor_address(self.input_name, int(d_input)) self.context.set_tensor_address(self.output_name, int(d_output)) self.context.execute_async_v3(stream.handle) cuda.memcpy_dtoh(output_buf, d_output) # 后处理 pred np.argmax(output_buf, axis1)[0] return pred这里我用 pycuda 简化了显存分配但生产环境你会希望把所有缓冲区的内存分配提前做好循环复用不要把cuda.mem_alloc放在每次推理里。每帧都分配显存性能开销非常大。预处理部分再次强调必须和训练脚本一致。很多同学部署时图省事用 OpenCV 读图后直接 resize 到 224x224既不先 resize 到 256 再中心裁剪也不做归一化最后精度跌得一塌糊涂我见过最夸张的案例是从 92% 掉到 70%排查了半天才发现是预处理的问题。4.3 性能验证与精度对比部署完成后一定要做两件验证性能验证和精度一致性验证。性能验证比较容易用 100 张图连续推理统计平均耗时和 P99 耗时。我这里提供一个简单的计时框架import time def benchmark(trt_infer, images, warmup10, repeat100): # warmup for img in images[:warmup]: trt_infer.infer(img) start time.time() for img in images[:repeat]: trt_infer.infer(img) avg_ms (time.time() - start) / repeat * 1000 return avg_ms精度一致性验证则要考虑 TensorRT 输出和 PyTorch 输出的差异来源。FP32 engine 的 softmax 输出和 PyTorch 的差异通常在 1e-5 到 1e-4 量级FP16 engine 的差异会到 1e-3 量级但 top-1 类别一般不会变。需要特别注意的是TensorRT 默认输出的是 logits不是 softmax 概率如果你在训练时习惯直接在模型里做 softmax导出 ONNX 时要把 softmax 也导进去否则部署端需要额外做一次 softmax。我自己偏好导出不包含 softmax 的裸 logits因为 softmax 是单调函数不影响 argmax 结果还能省一点计算。5. 常见问题与排查技巧实录5.1 TensorRT 常见报错与排查表部署过程中遇到的坑不少我整理成了速查表里面每一行都是实际踩过的报错或现象根本原因解决方案[TRT] Failed to parse ONNX modelONNX 里存在不支持的算子或版本不兼容降低 opset 到 12检查是否有自定义算子用网络层剥离法排查构建 engine 时显存不足workspace 设置过大或机器显存本身不足调低set_memory_pool_limit换更小的 batch增加 swapengine 加载报invalid engineengine 与 TensorRT 版本或 GPU 架构不匹配用目标环境重新构建 engine避免跨机器拷贝FP16 后精度显著下降模型对精度敏感某些层不适合 FP16关 FP16 用 FP32或做逐层精度分析对敏感层强制 FP32PyTorch 输出和 TensorRT 输出差异大预处理不一致导致输入分布不同逐项核对缩放、裁剪、归一化、通道顺序推理第一帧特别慢TensorRT 首次推理时做 context 初始化预热 5 到 10 次生产服务建议进程启动时先跑一次空推理Could not find cudnn报错CUDA/cuDNN 版本和 TensorRT 不匹配检查nvidia-smi与 TensorRT 官方兼容矩阵动态 batch 推理报 shape 不匹配context 没有设置 input shape每次推理前调用context.set_input_shape5.2 一次精度掉点的完整排查实录有一个印象很深的案例我这边模型 FP32 部署后单张推理性能完全达标但验证集 top-1 精度只有 89%训练时是 93%掉了 4 个点。当时第一反应是不是 TensorRT 转换出了问题但排除了算子兼容性问题又怀疑是量化误差但这个模型跑的是 FP32不该有这么大误差。后来逐项排查发现是预处理不一致。训练时用的是Resize(256)之后CenterCrop(224)但部署端前处理代码写成了直接cv2.resize(image, (224, 224))同时还把 BGR 通道忘了转成 RGB。这种问题在分类任务里特别隐蔽因为模型依然能输出一个看似合理的结果只是精度会莫名其妙掉几个点。排查这类问题我建议做一个“共享数据”的回归测试把同一张图分别喂给 PyTorch 模型和 TensorRT engine对比两者的 logits 分布。如果最大差异超过 1e-2基本可以断定不是 TensorRT 的问题而是前处理或输入缓冲区的问题。具体做法是固定输入是一张已知图片分别用训练时的预处理代码和部署端的预处理代码处理一遍打印出归一化后的像素均值、方差一眼就能看出差异。5.3 版本兼容性问题的经验总结关于 TensorRT 的版本兼容性我强调再多都不为过。TensorRT 8.4、8.5、8.6 三个版本之间Python API 都有一些小的行为差异特别是create_builder_config、set_memory_pool_limit这类 API8.5 之前叫set_max_workspace_size8.5 之后改成了 set memory pool 的写法直接抄网上老教程经常编译报错。我这里有一张实测过的兼容表仅供参考TensorRT 版本CUDA 版本PyTorch 推荐版本ONNX 推荐版本8.2.x11.31.10 - 1.121.10 - 1.128.4.x11.61.12 - 1.131.12 - 1.138.5.x11.7 / 11.81.13 - 2.01.138.6.x11.8 / 12.02.01.13另一个常见问题是 PyTorch 2.x 里torch.onnx.export的默认行为变了比如 PyTorch 2.6 开始torch.load的weights_only参数默认值改成了 True这跟 ONNX 导出本身没关系但如果你在部署脚本里加载历史权重就会莫名其妙报错。旧权重如果包含torch.save(model.state_dict())生成的普通 Map 对象加载时可能需要显式设置weights_onlyTrue或pickle_module等参数。建议所有部署相关脚本里统一写清楚加载参数避免环境升级带来的隐形 break。6. 部署在不同目标平台的补充说明6.1 服务器 GPU 与 Jetson 的差异同一套代码在服务器和 Jetson 上表现差异很大。我测试过 RTX 3060 和 Jetson Orin Nano 两个平台mobilenetv2 在 Orin Nano 上用 FP32 engine 推理大约 8 到 12 毫秒用 FP16 engine 可以压到 5 毫秒以内。关键是 Jetson 平台的 TensorRT 版本和服务器版本往往是分开发布的NVIDIA 官方提供 JetPack 里的 TensorRT 是定制过的直接用pip install tensorrt装到 Jetson 上很可能失败必须使用 JetPack 自带的版本。在 Jetson 上构建 engine 还容易出现一个问题显存有限构建过程中workspace设置太大会导致程序被 OOM kill。建议在 Jetson 上把 workspace 限制在 512MB 以内mobilenetv2 这种小模型实际用不到那么多。6.2 量化与精度校准的补充INT8 量化是进一步提升性能的手段但 mobilenetv2 这种轻量模型对量化比较敏感直接 PTQPost-Training Quantization掉点可能比较明显。如果你一定要做 INT8建议做好校正集和逐层精度分析至少保证校正集包含足够多样的真实数据而不是随便拍几张图片凑数。我的实测结果是mobilenetv2 在 RTX 3060 上 INT8 engine 比 FP16 大概再快 20% 到 30%但精度从 93% 掉到 90.5% 左右对分类任务来说还能接受对更细粒度的任务则不建议冒险。7. 实际部署中的几个技巧与体会最后分享几个我长期实践下来觉得很实用的技巧。加载 engine 文件时建议把 engine 序列化到本地文件后用一个统一的 EngineLoader 类管理加载逻辑代码里不要出现裸的open()read()。原因是 engine 是二进制协议格式加载失败时报错信息往往很难看统一封装后可以捕获异常、打印清晰的错误提示排查问题效率高很多。推理循环里输入图像的 buffer 复用非常重要。生产环境的高吞吐推理瓶颈往往不在 kernel 执行而在 Host 与 Device 之间的内存拷贝。我建议用cuda-python或pycuda提前分配好固定的 device buffer每次推理只更新 buffer 内容不要反复mem_alloc和free。实测下来这个改动配合 batch4 推理吞吐能提升接近一倍。batch 策略上如果推理服务需要支持高并发可以尝试在服务端攒 batch同一时刻到达的请求组合成一个 batch 输入 TensorRT而不是每个请求单独推理。mobilenetv2 模型小batch8 时 GPU 还有很充裕的并行度吞吐提升非常可观。但积攒 batch 会带来额外的延迟需要根据流量模型权衡。模型更新和热加载上建议引擎文件按版本命名比如mobilenetv2_v1.0_fp16.engine服务启动时通过配置指定引擎路径新版本模型文件放好后做一次平滑重载。TensorRT engine 本身是二进制文件升级模型不需要重新编译服务。我个人的体会是从 PyTorch 训练到 TensorRT 部署整条链路真正花时间的往往不是模型代码而是这些版本的匹配、预处理的对齐、缓冲区的管理这类工程细节。把这些细节做扎实mobilenetv2 的分类部署其实一两天就能全部搞定。后面你如果要迁移到其他模型比如 YOLO 系列或者更重的分类网络这套方法论完全可以直接搬过去变的只是网络结构解析那一步部署架构和排查思路都是一样的。本文还有配套的精品资源点击获取