ONNXRuntime中BERT模型导出为ONNX格式输出NaN问题的分析与解决

问题背景

在使用ONNXRuntime进行BERT模型推理时，开发者经常会遇到模型输出为NaN（Not a Number）的问题。这种情况通常发生在将Hugging Face Transformers中的BERT模型导出为ONNX格式后，使用C++接口进行推理时。本文将从技术角度深入分析这一问题的成因，并提供有效的解决方案。

问题现象

开发者报告了以下典型现象：

1. 使用Python的transformers.onnx工具导出BERT模型到ONNX格式
2. 使用ONNX Runtime的C++接口进行推理时，CPU上输出为"-nan"，GPU上输出为"nan"
3. 该问题在ONNX Runtime的多个版本（1.11.1至1.17.1）中均存在

根本原因分析

经过技术分析，导致这一问题的可能原因包括：

1. 输入张量顺序错误：在导出ONNX模型时，attention_mask可能与其他输入张量顺序错位。ONNX模型对输入顺序有严格要求，顺序不匹配会导致计算异常。

2. 模型导出工具版本问题：早期版本的transformers.onnx工具可能存在导出逻辑缺陷，导致生成的ONNX模型结构不完整或参数异常。

3. 数值稳定性问题：BERT模型中的softmax或layer normalization等操作在特定输入下可能导致数值不稳定，产生NaN。

4. 执行提供程序兼容性问题：不同版本的ONNX Runtime对CUDA执行提供程序的支持可能存在差异。

解决方案

1. 验证ONNX模型结构

使用Netron工具可视化检查导出的ONNX模型：

• 确认输入节点名称和顺序是否符合预期
• 检查模型中的attention_mask输入位置是否正确
• 验证模型各层参数是否完整

2. 升级相关工具版本

建议采取以下升级措施：

• 将Hugging Face Transformers升级到最新稳定版
• 使用ONNX Runtime 1.21或更高版本
• 确保CUDA版本与ONNX Runtime兼容（建议CUDA 11.x）

3. 检查输入数据

确保推理时的输入数据：

• 数据类型与模型预期一致（通常是float32或int64）
• 数值范围合理（避免极端值导致数值不稳定）
• 张量形状与模型输入要求匹配

4. 使用官方验证脚本

ONNX Runtime提供了BERT模型的验证脚本，可以用来确认模型导出和推理的正确性：

python -m onnxruntime.transformers.models.bert.eval_squad

最佳实践建议

1. 标准化导出流程：建立统一的模型导出流程，记录使用的工具版本和参数。

2. 版本控制：对ONNX模型文件进行版本管理，记录导出环境和参数。

3. 渐进式验证：从简单输入开始逐步验证模型，先确保小批量数据能正确推理。

4. 性能监控：实现推理过程的数值监控，及时发现异常值。

5. 跨平台测试：在CPU和GPU环境下分别测试模型，确保兼容性。

总结

BERT模型导出为ONNX格式后出现NaN输出是一个典型的技术问题，通常与模型导出过程或推理环境配置有关。通过升级工具版本、仔细验证模型结构和输入数据，大多数情况下可以解决这一问题。开发者应当建立规范的模型导出和验证流程，以确保深度学习模型在不同平台间的顺利迁移和部署。

随着ONNX生态的不断成熟，这类问题的发生频率正在降低，但保持工具链的更新和标准化操作流程仍然是预防问题的有效手段。