4. 工业AI模型部署流程

工业AI项目的落地难点往往不在训练阶段，而在“训练好的模型如何稳定、高效地跑在IGX边缘设备上”。一个在服务器上精度达标的PyTorch模型，直接搬到IGX Orin上可能面临算子不支持、显存溢出、推理延迟超标、动态输入尺寸报错等一系列问题。本节围绕“训练→ONNX→TensorRT引擎→部署验证”这条主线，讲清每一步的实际操作、关键参数与判断依据。

一、为什么需要迁移链路：从框架模型到推理引擎

训练框架（PyTorch/TensorFlow）为了支持反向传播和灵活调试，保留了完整的计算图和大量动态特性；而边缘推理追求的是低延迟、低功耗、确定性。二者目标不同，因此需要一条“翻译+编译+优化”的链路：先把框架模型导出为与框架解耦的中间表示（ONNX），再由TensorRT针对IGX的GPU架构（如Orin的Ampere架构）做算子融合、精度校准和内核自动调优，最终生成序列化的推理引擎（.plan文件）。

这条链路的核心价值在于：ONNX解决“跨框架可移植”，TensorRT解决“针对硬件极致优化”。理解这一点，就能明白为什么不能跳过ONNX直接部署，也能理解精度校准和动态形状为什么必须在这一层处理。

二、ONNX导出：把训练模型“翻译”成通用格式

以PyTorch为例，导出ONNX的核心是torch.onnx.export()。关键参数有三个：opset_version建议设为17或更高（对动态形状和较新算子支持更好）；input_names和output_names要显式命名，便于后续TensorRT绑定；dynamic_axes用于声明哪些维度是动态的。例如一个输入为[N,3,H,W]的检测模型，若希望batch和分辨率都可变，应声明：

dynamic_axes = {
    'input':  {0: 'batch', 2: 'height', 3: 'width'},
    'output': {0: 'batch'}
}
torch.onnx.export(model, dummy_input, 'model.onnx',
                  opset_version=17,
                  input_names=['input'],
                  output_names=['output'],
                  dynamic_axes=dynamic_axes)

导出后必须做两件事：一是用onnx.checker.check_model()校验模型结构完整性；二是用ONNX Runtime跑一遍推理，与PyTorch原始输出做数值比对。若最大绝对误差超过1e-4，说明导出过程中有算子被错误转换，需要回查。常见问题包括：自定义算子无对应ONNX实现、控制流（if/loop）导出失败、动态轴命名冲突。此时可考虑将自定义算子拆解为ONNX原生算子组合，或使用torch.onnx.register_custom_op_symbolic注册符号函数。

三、TensorRT引擎构建：算子融合与精度选择

拿到ONNX后，用trtexec或TensorRT Builder API构建引擎。以trtexec为例：

trtexec --onnx=model.onnx \
        --saveEngine=model_fp16.plan \
        --fp16 \
        --minShapes=input:1x3x320x320 \
        --optShapes=input:1x3x640x640 \
        --maxShapes=input:1x3x1280x1280 \
        --workspace=4096

这里--fp16启用半精度，通常能带来1.5~2倍加速且精度损失极小；--minShapes/optShapes/maxShapes定义动态形状的三个档位，其中optShapes最重要——TensorRT会针对该尺寸选择最优内核并分配显存，实际推理尺寸越接近opt，性能越好。--workspace指定构建时可用的显存上限（单位MB），过小会导致某些层无法使用最优算法。

若需要INT8推理，则必须提供校准集。校准的本质是：用一批有代表性的输入数据统计各层激活值的动态范围，生成缩放因子。校准集应覆盖实际部署中可能出现的各种场景（不同光照、角度、目标密度），数量500~1000张即可。使用IInt8EntropyCalibrator2时，需实现getBatch()方法按批次提供数据。校准完成后，务必用验证集比对INT8引擎与FP32/FP16引擎的输出，若mAP下降超过1个百分点，应检查校准集是否偏斜，或对敏感层回退到FP16。

四、动态形状支持：Profile配置与运行时绑定

工业场景中，输入尺寸常随相机分辨率或ROI裁剪而变化。TensorRT通过Optimization Profile支持动态形状。一个引擎可绑定多个Profile，但同一时刻只有一个生效。运行时需调用setBindingDimensions()设置实际输入尺寸，并确保该尺寸落在某个Profile的[min, max]区间内。若超出，推理会直接失败。

实践中有两个易错点：其一，optShapes应设为实际最常用的尺寸，而非简单取min和max的中间值；其二，若模型包含reshape或transpose等对形状敏感的算子，动态维度变化可能导致引擎内部张量维度推导错误，此时需在ONNX导出时固定部分维度，或改用多个静态引擎分别部署。

五、部署验证：从“能跑”到“跑得好”

引擎构建完成只是开始，部署验证才是确保工业现场稳定运行的关键。验证应覆盖四个维度：

• 精度一致性：用同一批测试样本，分别跑ONNX Runtime和TensorRT引擎，逐输出比对余弦相似度，要求大于0.999；对检测模型还需比对mAP。
• 延迟与吞吐：用trtexec --loadEngine=model.plan --shapes=input:1x3x640x640测量单流延迟；再用多流并发测吞吐。关注P50和P99延迟，P99过高说明存在调度抖动。
• 显存与功耗：通过nvidia-smi或tegrastats监控推理过程中的显存占用和功耗，确保不超过IGX模组的热设计功耗（TDP）预算。
• 长时间稳定性：连续运行24小时，观察显存是否持续增长（内存泄漏）、延迟是否漂移。工业设备通常要求7×24小时运行，这一步不可省略。

最后需要强调：整条链路中，ONNX是“合同”，TensorRT引擎是“成品”，而部署验证是“验收”。任何一步偷懒，都可能在现场以花屏、漏检、宕机等形式暴露出来。把每一步的输入、输出和判断标准都明确下来，工业AI模型才能真正从实验室走进产线。