视觉模型部署实战：从 PyTorch 到 TensorRT 的性能跃迁
## 训练与部署是两回事
很多算法工程师能训练出精度不错的模型，但一到部署就遇到问题：单张推理耗时 200ms，无法满足实时要求；或者显存占用过高，单卡只能跑一路视频流。
差距来自一个事实：**PyTorch 是为训练设计的，不是为推理设计的**。
## 完整部署路径
```
PyTorch (.pt)
↓ 导出
ONNX (.onnx)
↓ 优化 + 转换
TensorRT Engine (.plan)
↓ 部署
推理服务
```
## 第一步：PyTorch → ONNX
```python
import torch
model.eval()
dummy_input = torch.randn(1, 3, 640, 640).cuda()
torch.onnx.export(
model,
dummy_input,
"model.onnx",
opset_version=17,
input_names=["images"],
output_names=["output"],
dynamic_axes={
"images": {0: "batch"},
"output": {0: "batch"}
}
)
```
**常见坑**：
- 动态 shape 支持需要正确配置 `dynamic_axes`
- 某些算子在低版本 opset 不支持，需提升版本
- 导出后必须验证输出一致性（与 PyTorch 对比，误差 < 1e-3）
## 第二步：ONNX 优化
```python
import onnx
from onnxruntime.transformers import optimizer
# 图优化：算子融合、常量折叠、消除冗余节点
optimized = optimizer.optimize_model(
"model.onnx",
model_type="bert", # 或自定义
num_heads=0,
hidden_size=0,
)
optimized.save_model_to_file("model_opt.onnx")
```
同时可用 ONNX Runtime 做推理测试，确认精度无损后再进入下一步。
## 第三步：ONNX → TensorRT
### FP16 量化（最常用）
```bash
trtexec --onnx=model_opt.onnx \
--saveEngine=model_fp16.plan \
--fp16 \
--workspace=4096 \
--verbose
```
FP16 通常带来 1.8-2.5 倍加速，精度损失可忽略（< 0.5% mAP）。
### INT8 量化（激进方案）
INT8 需要校准数据集来确定激活值的量化范围：
```python
from polygraphy.backend.trt import Calibrator, CreateConfig
calibrator = Calibrator(
data_loader=calibration_loader, # 500-1000 张代表性样本
cache="calib.cache"
)
config = CreateConfig(int8=True, calibrator=calibrator)
```
INT8 可再提升 1.5-2 倍，但精度损失达 1-3% mAP，需谨慎评估。
## 实测性能对比
以 YOLOv8s 在 640×640 输入、RTX 4090 上的实测数据为例：
| 格式 | 单张延迟 | 吞吐（FPS） | 显存占用 | mAP 损失 |
|------|---------|------------|---------|---------|
| PyTorch FP32 | 12.4 ms | 81 | 1.8 GB | 基准 |
| ONNX Runtime | 8.6 ms | 116 | 1.5 GB | 0 |
| TensorRT FP16 | 3.5 ms | 286 | 1.1 GB | < 0.3% |
| TensorRT INT8 | 2.6 ms | 385 | 0.9 GB | ~1.5% |
从 PyTorch 到 TensorRT FP16，**吞吐提升 3.5 倍**；若再叠加批处理与流水线优化，端到端可达 4.7 倍。
## 视频流部署的关键设计
单张推理快不等于系统吞吐高。多路视频流场景还需要：
1. **解码与推理流水线**：用硬解码（NVDEC）+ 异步推理，避免串行等待
2. **动态批处理**：累积多路帧组成 batch 一次推理
3. **跳帧策略**：非关键场景降低采样帧率
4. **多进程隔离**：按 GPU 划分进程，避免 GIL 与显存竞争
典型配置下，单张 4090 可支持 8-16 路 1080p 视频的实时分析。
## 落地建议
部署优化应该**前置到模型设计阶段**考虑：选择对量化友好的网络结构、避免使用不支持的自定义算子、控制输入分辨率。事后优化永远比事前设计成本更高。
启数AI教育 AI 视觉课程第四阶段以 TensorRT 部署为核心，要求学员完成从训练到边缘设备部署的完整链路。