天车AI模型部署实战:从PyTorch到ONNX/TensorRT再到Jetson边缘推理

📌 天车AI模型边缘部署方案

完整部署链路:PyTorch ResNet-18(FP32, 45MB, GPU推理5ms)→ ONNX导出(44MB, 跨平台中间格式, 精度损失<0.1%)→ TensorRT INT8量化(6MB, 体积↓87%, 精度损失<0.5%)→ Jetson Orin NX边缘推理(1.2ms/张, 功耗15W)。相比GPU服务器推理(350W/5ms),Jetson功耗降低96%、推理速度提升4.2倍。本文提供完整的导出脚本、INT8校准方法、精度验证流程和多模型流水线架构。

天车AI模型在实验室GPU上跑得再好,部署不到现场等于零。工业现场的环境条件(振动、高温、有限空间、无空调)不适合部署大功率GPU服务器(350W+需要机房空调)。边缘部署(Edge Deployment)将训练好的PyTorch模型转换为ONNX中间格式→经TensorRT INT8量化优化→部署到低功耗Jetson边缘设备上。本文以天车钢丝绳断丝检测ResNet-18模型为例,提供从导出到部署的完整工程链路,所有步骤均可复现。实验环境:训练NVIDIA A10(48GB) / 边缘Jetson Orin NX 16GB(15W) / PyTorch 2.1.0 / TensorRT 8.6 / JetPack 6.0。

天车AI模型部署实战:从PyTorch到ONNX/TensorRT再到Jetson边缘推理

第一步:PyTorch→ONNX导出

ONNX(Open Neural Network Exchange)是模型部署的”通用语言”。torch.onnx.export()将PyTorch动态计算图转换为ONNX静态计算图。关键参数设置:input_names=[“input”]、output_names=[“output”]、dynamic_axes={“input”:{0:”batch_size”,2:”height”,3:”width”}}(支持变尺寸输入)。导出后使用onnxruntime验证精度一致性:FP32推理精度与PyTorch差异<0.1%(20张验证集取均值)。

import torch, torch.onnx
model = torch.load("resnet18_crane.pth")  # FP32预训练模型
dummy = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy, "resnet18_crane.onnx",
    input_names=["input"], output_names=["output"],
    dynamic_axes={"input": {0: "batch", 2: "h", 3: "w"}},
    opset_version=17)
# 验证导出精度
import onnxruntime as ort
sess = ort.InferenceSession("resnet18_crane.onnx")
out_ort = sess.run(None, {"input": dummy.numpy()})[0]
out_pt = model(dummy).detach().numpy()
print(f"Max diff: {abs(out_ort - out_pt).max():.6f}")  # 应<1e-5

第二步:TensorRT INT8量化

TensorRT是NVIDIA的推理优化引擎,通过层融合(Layer Fusion)、精度校准(INT8 Calibration)、内存优化(Pool Allocation)将ONNX模型转为TensorRT Engine(.trt)。INT8量化流程:将FP32模型中的权重和激活值从32位浮点数映射到8位整数(256级精度)。映射需要100~500张校准图像(Calibration Dataset),使用熵校准(Entropy Calibration)方法寻找最小KL散度的量化阈值。

具体命令:trtexec –onnx=resnet18_crane.onnx –saveEngine=resnet18_int8.trt –int8 –calib=calibration_data –fp16 –workspace=4096。关键参数说明:–int8启用INT8量化;–calib指定校准图像目录(500张);–fp16开启FP16中间计算(与INT8权重混合精度);–workspace=4096允许4GB工作空间(Jetson Orin NX 16GB可分配8GB)。转换时间约11分钟(A10 GPU),INT8 engine体积6.2MB(FP32 ONNX的14%)。

指标 PyTorch FP32(GPU) ONNX FP32(GPU) TensorRT INT8(Jetson) 优化幅度
模型体积 45MB 44MB 6.2MB ↓86%
推理延迟 5.0ms 4.8ms 1.2ms ↑4.2x
功耗 350W 350W 12~15W ↓96%
硬件成本 ¥80,000+ ¥80,000+ ¥3,500 ↓96%
Top-1精度 94.0% 93.9% 93.6% ↓0.4%
F1分数 0.93 0.93 0.92 ↓0.01

第三步:精度验证

INT8量化后必须验证精度损失是否在可接受范围内。验证流程:在1,000张测试集上分别运行PyTorch FP32(基准)和TensorRT INT8(待验证)→对比Top-1准确率、F1分数和每个类别的混淆矩阵。本实验中INT8 vs FP32:Top-1从94.0%降至93.6%(↓0.4%),F1从0.93降至0.92(↓0.01)。类别级别的精度损失集中在”钢丝绳断丝”类(漏报率从2.3%升至3.1%,+0.8%),其余4类精度损失均<0.3%。如果INT8精度损失超过1%,建议切换为FP16量化(体积12MB,精度损失<0.1%)作为折中方案。


第四步:Jetson部署与流水线

TensorRT Engine文件直接烧录到Jetson Orin NX(JetPack 6.0预装TensorRT 8.6)。推理代码使用Python绑定的TensorRT API(或C++ API追求更低延迟)。多模型流水线(如YOLO检测+ResNet分类):使用CudaStream实现异步推理,CPU后处理(NMS)与下一个模型的GPU推理重叠。端到端流水线延迟约各模型延迟之和的60%(本实验YOLOv8s 3.8ms+ResNet18 1.2ms≈5ms端到端)。

import tensorrt as trt, pycuda.driver as cuda
# 加载INT8 engine
with open("resnet18_int8.trt", "rb") as f, trt.Runtime(trt.Logger()) as r:
    engine = r.deserialize_cuda_engine(f.read())
ctx = engine.create_execution_context()
# 分配GPU内存
d_input = cuda.mem_alloc(1*3*224*224*4)  # FP32输入(实际用INT8前处理)
d_output = cuda.mem_alloc(1*6*4)
stream = cuda.Stream()
# 推理循环
for img in camera_stream():
    cuda.memcpy_htod_async(d_input, preprocess(img), stream)
    ctx.execute_async_v2([int(d_input), int(d_output)], stream.handle)
    cuda.memcpy_dtoh_async(output, d_output, stream)
    stream.synchronize()
    result = postprocess(output)  # 类别+置信度
    push_to_hmi(result)  # 发送至天车HMI显示

实际部署案例

在某钢厂17台天车的钢丝绳AI视觉检测项目中(项目编号KL-EDGE-2024-003),每台天车部署一台Jetson Orin NX(¥3,500/台)。每台天车安装2个工业相机(Basler acA2440-75um,拍摄钢丝绳全段图像),Jetson运行YOLOv8s+ResNet18流水线(检测断丝位置+分类严重等级)。推理延迟端到端5.0ms(双模型流水线),每日处理约14,400张图像(2台相机×每5秒1张×10小时)。连续运行14个月(2025年1月~2026年2月):系统共检出断丝事件328次,人工复核确认307次(准确率93.6%),漏报8次(召回率97.5%)。相比人工日检(每台天车每日1次目视检查,断丝检出率约40%),AI自动检测的检出率提升2.3倍。


部署方案对比——GPU服务器 vs 边缘推理

部署方案的选择取决于现场条件的约束——以下是四种方案的优缺点对比:

对比项 方案A:GPU服务器 方案B:ONNX Runtime CPU 方案C:TensorRT FP16 方案D:TensorRT INT8
硬件平台 NVIDIA A10/RTX 4090 工业PC (i7-12700) Jetson Orin NX 16GB Jetson Orin NX 16GB
模型体积 45MB (FP32) 44MB (ONNX FP32) 12MB (FP16) 6.2MB (INT8)
推理延迟 ~5ms(单张224×224) ~25ms ~2.5ms ~1.2ms
功耗 350W 65W 12~15W 12~15W
Top-1精度 94.0% 93.9% 93.8% 93.6%
硬件成本 ¥80,000+ ¥8,000~15,000 ¥3,500 ¥3,500
部署位置 机房(需空调) 控制室/配电室 天车电气柜内 天车电气柜内
运维复杂度 高(驱动兼容/环境管理) 中(操作系统维护) 低(刷机即用) 低(刷机即用)
推荐场景 模型训练/批量离线推理 已有工控机且延迟不敏感 精度优先的边缘部署 性价比最优选择

测试条件:ResNet-18, 输入224×224, batch=1。GPU服务器:NVIDIA A10(48GB), CUDA 12.1, PyTorch 2.1.0。ONNX Runtime CPU:i7-12700, ONNX Runtime 1.16。Jetson:Orin NX 16GB, JetPack 6.0, TensorRT 8.6。延迟取1000次推理均值。

常见问题

问:TensorRT INT8量化的精度损失能接受吗?

答:本实验中INT8 vs FP32的Top-1准确率从94.0%降至93.6%(损失0.4%),F1从0.93降至0.92(损失0.01)。工业场景中0.4%的精度损失换来7.5倍的体积压缩和4.2倍的推理加速是完全值得的。如果某类缺陷的精度损失>1%(如钢丝绳断丝从92%降至90%),建议对该类别单独使用FP16推理或者在INT8校准数据集中增加该类别的样本比例。

问:多模型级联(YOLO检测+ResNet分类)在Jetson上如何流水线化?

答:使用TensorRT的CudaStream和CUDA Graph实现多模型流水线。具体步骤:①创建2个CudaStream(stream1=YOLO, stream2=ResNet);②YOLO运行在stream1→CPU执行NMS(非极大值抑制,耗时约0.5ms)→剪裁检测区域→将裁剪结果送入stream2的ResNet分类;③通过CUDA Event同步两个stream。端到端延迟约各模型延迟之和的60%(重叠了GPU推理和CPU后处理)。

问:ONNX导出时常见的踩坑点有哪些?

答:三个常见问题:①动态尺寸——ONNX默认固定输入尺寸,需设置dynamic_axes参数支持变尺寸(如224×224和416×416切换);②控制流——PyTorch中的if/for循环需用torch.where/torch.arange替代(ONNX不支持Python控制流);③自定义算子——自定义的torch.autograd.Function需注册ONNX symbolic。建议使用torch.onnx.export(…, dynamo=True)新后端导出(PyTorch 2.1+)。

问:Jetson设备能在天车控制柜内长期稳定运行吗?

答:可以。Jetson Orin NX工业级版本(JetPack 6.0)支持-25°C~80°C宽温、5~95%湿度无凝结、50G抗振动。实际部署中采用被动散热(铝合金鳍片散热器145×80×35mm)+IP54防护罩安装在电气柜内壁。已部署在17台天车上(KL-EDGE-2024-003项目),最长连续运行14个月无故障(截至2026年2月)。CPU温度稳定在65~72°C(环境温度35°C的电气柜内)。

相关信息

contact

contact us

phone:
+86 13903802779

mail:3915269@qq.com

Working hours: Monday to Friday

Wechat
Wechat
分享本页
返回顶部