TensorRT中FP16精度转换导致NaN输出的问题分析与解决方案

问题背景

在深度学习模型部署过程中，将模型从FP32精度转换为FP16精度是常见的优化手段，可以显著提升推理速度并减少显存占用。然而，在使用TensorRT进行FP16精度转换时，开发者可能会遇到模型输出NaN（非数字）的问题。

问题现象

当尝试将一个基于NAFNet架构的模型从FP32转换为FP16精度时，模型在FP32精度下能够产生正确的输出，但在FP16精度下却输出NaN值。通过Polygraphy工具比较ONNX Runtime和TensorRT的输出时，发现TRT-FP16的输出无效。

问题分析

经过深入分析，发现这个问题可能由以下几个原因导致：

1. FP16精度范围限制：FP16的数值范围（约±65504）远小于FP32（约±3.4×10³⁸），在模型计算过程中容易出现数值溢出。

2. 特定操作的影响：模型中的CumSum（累积求和）操作特别容易导致数值超出FP16的范围，因为随着求和次数的增加，数值会不断累积增大。

3. TensorRT版本差异：有开发者报告称，在TensorRT 8.5版本中相同的模型可以正常工作，但在10.6和10.7版本中会出现NaN输出，这表明TensorRT内部实现的变化可能影响了FP16的数值稳定性。

解决方案

1. 混合精度策略

通过分析模型结构，可以识别出对数值精度敏感的关键层，并强制这些层保持FP32精度：

import onnx
model_path = 'model.onnx'
model = onnx.load(model_path)
graph = model.graph
matching_layer = []
for node in graph.node:
    for output in node.output:
        if 'sca' in output:  # 识别所有包含'sca'的层
            matching_layer.append(node)

layer_prec_str = ''
for layer in matching_layer:
    layer_prec_str += layer.name + ':fp32,'
layer_prec_str = layer_prec_str.removesuffix(',')

然后使用trtexec工具构建混合精度引擎：

trtexec --onnx=model.onnx --saveEngine=model_fp16.plan --fp16 \
        --precisionConstraints=prefer --layerPrecisions=<上面生成的层精度字符串>

2. 使用Polygraphy调试工具

TensorRT提供的Polygraphy工具可以帮助开发者调试精度问题：

polygraphy debug precision model.onnx --fp16 --tactic-sources cublas \
           --verbose -p float32 \
           --check polygraphy run polygraphy_debug.engine --trt \
           --load-inputs input.json --load-outputs golden.json --abs 1e-2

这个命令会尝试找出导致数值问题的层，并自动调整这些层的精度。

3. 考虑使用ONNX Runtime

如果模型在ONNX Runtime中能够正确运行FP16推理，而性能差异可以接受，可以考虑直接使用ONNX Runtime作为推理后端。

最佳实践建议

1. 逐步转换策略：不要一次性将整个模型转换为FP16，而是逐步转换并验证每一部分的输出。

2. 数值范围检查：在模型训练和转换过程中，添加数值范围检查机制，识别潜在的数值不稳定操作。

3. 版本兼容性测试：如果可能，在不同版本的TensorRT上测试模型，选择最稳定的版本。

4. 模型结构调整：对于特别容易出现数值问题的结构（如深层的累积操作），考虑从模型设计层面进行优化。

结论

FP16精度转换虽然能带来性能提升，但也引入了数值稳定性挑战。通过混合精度策略和适当的工具链支持，开发者可以在保持模型精度的同时获得FP16的性能优势。理解模型的计算特性和TensorRT的工作原理是解决这类问题的关键。