TensorRT模型输出NaN问题的分析与解决

问题背景

在使用TensorRT部署SuperGlue模型时，开发者遇到了一个常见但棘手的问题：模型推理结果总是返回NaN（非数字）值。这种情况通常发生在模型转换或推理过程中出现数值不稳定的情况下。

环境配置

该问题出现在以下环境中：

• TensorRT版本：8.6.2.3
• GPU设备：NVIDIA Jetson Orin NX（16GB内存）
• CUDA版本：12.02.140
• cuDNN版本：8.9.4.25
• 操作系统：Ubuntu 22.04 LTS
• Python版本：3.10.12

问题现象

开发者成功将SuperGlue模型转换为ONNX格式，并进一步转换为TensorRT引擎。模型能够正常运行且不报错，但输出结果始终为NaN。当开发者尝试设置输出形状大于预期形状时，在预期形状部分得到NaN，其他部分则得到0.0。

根本原因分析

经过技术社区的分析，问题根源在于TensorRT的FP16（半精度浮点）优化标志。在模型转换过程中，开发者启用了FP16模式（通过config.set_flag(trt.BuilderFlag.FP16)），这可能导致数值精度不足，从而引发数值下溢或上溢问题。

解决方案

解决此问题的方法非常简单但有效：

1. 禁用FP16优化：在构建TensorRT引擎时，注释掉或移除设置FP16标志的代码行。

# 注释掉这行代码
# config.set_flag(trt.BuilderFlag.FP16)

2. 使用FP32精度：保持默认的FP32（单精度浮点）精度可以避免因精度不足导致的数值不稳定问题。

技术原理

FP16（半精度浮点）使用16位表示浮点数，相比FP32（32位）可以：

• 减少内存占用
• 提高计算速度
• 增加吞吐量

但同时也会带来：

• 数值表示范围缩小
• 精度降低
• 更容易出现数值不稳定

在SuperGlue这类计算机视觉模型中，某些层的计算可能对数值精度特别敏感，使用FP16可能导致中间结果超出表示范围或损失关键精度信息，最终导致输出为NaN。

最佳实践建议

1. 精度选择策略：

   • 首先尝试FP32模式确保模型正确性
   • 在验证模型稳定后，可尝试启用FP16进行性能优化
   • 对于特别敏感的模型，考虑混合精度策略

2. 数值稳定性检查：

   • 在模型转换过程中添加数值范围检查
   • 监控中间结果的数值范围
   • 使用TensorRT的调试工具分析各层输出

3. 模型优化：

   • 对于FP16模式，考虑添加数值稳定层
   • 调整模型参数范围以适应FP16表示
   • 使用适当的归一化技术

总结

TensorRT模型输出NaN是一个常见问题，通常与数值精度设置有关。通过禁用FP16优化可以解决大多数此类问题。在实际部署中，开发者需要在模型精度和推理性能之间找到平衡点，根据具体应用场景选择适当的精度级别。