trtexec执行模型时显存不足如何优化？

Question

在使用 `trtexec` 执行 TensorRT 模型时，常遇到显存不足（Out of Memory, OOM）问题。这通常发生在模型规模大、输入尺寸高或批处理数据过多时。如何在有限显存下成功运行模型，是优化的关键。常见解决方法包括：降低输入 batch size、减小输入分辨率、启用 FP16 或 INT8 精度推理、使用动态形状时限制最大尺寸、拆分模型执行等。此外，TensorRT 提供的内存优化选项如 `--workspace` 和 `--minTiming` 也可辅助调整。如何结合具体场景选择合适的优化策略，是提升模型推理效率和稳定性的核心挑战。

Answer 1

一、问题背景与显存不足的常见原因

在使用 trtexec 工具执行 TensorRT 模型时，显存不足（OOM）是一个常见的瓶颈。尤其当模型结构复杂、输入尺寸大或批处理数据量多时，GPU 显存容易耗尽，导致推理失败。造成 OOM 的主要原因包括：

• 模型参数规模过大，如 ResNet-152 或 BERT 类模型。
• 高分辨率输入图像或视频帧。
• 批量（batch size）设置过高。
• 未启用低精度计算（FP16/INT8）。
• 动态形状配置不当，最大维度超出显存容量。

二、优化策略分析与优先级排序

针对上述问题，TensorRT 提供了多种内存优化手段。以下为常见策略及其优先级建议：

1. 降低 batch size：最直接有效的方式，减少每轮计算的数据量。
2. 减小输入分辨率：适用于视觉任务，如图像分类或目标检测。
3. 启用 FP16 或 INT8 推理：显著降低内存占用和提升推理速度。
4. 限制动态形状的最大尺寸：避免运行时分配过多显存。
5. 拆分模型执行：将模型切分为多个子图分别执行。
6. 调整 TensorRT 参数：如 --workspace 和 --minTiming 等。

三、trtexec 常用优化命令示例

优化选项	作用说明	示例命令片段
--batch	设置输入批大小	--batch=16
--fp16	启用 FP16 混合精度	--fp16
--int8	启用 INT8 量化	--int8 --calib=calibration.cache
--workspace	设置最大工作空间大小（单位 MB）	--workspace=512
--shapes	指定动态形状输入尺寸	--shapes=input:0x3x224x224
--minTiming	减少用于性能评估的时间测量次数	--minTiming=1

四、显存优化流程图

graph TD
A[开始] --> B{是否出现OOM?}
B -- 否 --> C[完成]
B -- 是 --> D[尝试降低batch size]
D --> E{是否仍OOM?}
E -- 否 --> F[完成]
E -- 是 --> G[尝试降低输入分辨率]
G --> H{是否仍OOM?}
H -- 否 --> I[完成]
H -- 是 --> J[启用FP16或INT8]
J --> K{是否仍OOM?}
K -- 否 --> L[完成]
K -- 是 --> M[限制动态形状最大尺寸]
M --> N{是否仍OOM?}
N -- 否 --> O[完成]
N -- 是 --> P[拆分模型或调整workspace]
P --> Q{是否仍OOM?}
Q -- 否 --> R[完成]
Q -- 是 --> S[考虑硬件升级]

五、结合具体场景选择合适策略

不同应用场景对性能与精度的要求不同，应灵活选择优化策略：

• 边缘设备部署：优先使用 INT8 + 小 batch + 输入降分辨率。
• 服务器端推理：可适当增加 batch size，但需监控显存使用情况。
• 训练后量化模型：应确保校准集准确，并验证精度无明显下降。
• 动态形状模型：合理设置 min/opt/max 形状，防止过度预留显存。