Triton推理服务器中TensorRT模型输入数据类型TYPE_INT64的支持问题解析

背景介绍

在NVIDIA Triton推理服务器的最新版本中，用户在使用TensorRT(TRT)模型时遇到了一个关于输入数据类型的兼容性问题。具体表现为：当模型配置文件(config.pbtxt)中指定输入数据类型为TYPE_INT32时，实际传入的TYPE_INT64类型数据会导致推理失败。

问题现象

用户在使用Triton服务器加载TensorRT模型时，遇到了以下错误信息：

ERROR: Failed to create instance: unexpected datatype TYPE_INT64 for inference input 'ref_seq', expecting TYPE_INT32 for t2s_encoder_0_0

查看模型配置文件，可以看到该模型明确指定了输入数据类型为TYPE_INT32：

input [
   {
     name: "ref_seq"
     data_type: TYPE_INT32
     dims: [ 1, -1 ]
   },
   ...
]

解决方案

通过将模型配置文件中的数据类型从TYPE_INT32修改为TYPE_INT64，问题得到解决。这一修改反映了TensorRT 10.0及以上版本对64位整数类型的支持能力。

技术原理深度解析

1. TensorRT版本演进：

   • 在TensorRT 10.0之前的版本中，对整数类型的支持主要限于32位(INT32)
   • 从TensorRT 10.0开始，正式引入了对64位整数(INT64)的完整支持

2. Triton服务器兼容性：

   • Triton 24.05是首个集成TensorRT 10.0的版本
   • 新版本能够正确识别和处理模型配置中的TYPE_INT64数据类型定义

3. 数据类型选择建议：

   • 对于需要大整数范围的场景，应优先使用INT64
   • 在内存敏感或性能关键场景中，可考虑使用INT32以节省资源
   • 必须确保模型配置文件中的数据类型定义与实际输入数据完全一致

最佳实践

1. 版本适配检查：

   • 在使用INT64数据类型前，确认TensorRT版本≥10.0
   • 确认Triton服务器版本≥24.05

2. 模型配置规范：

   • 明确定义所有输入输出的数据类型
   • 保持配置文件与实际数据的一致性
   • 对于动态维度，使用-1标记但要确保运行时维度合理

3. 性能考量：

   • INT64操作通常比INT32消耗更多内存和计算资源
   • 在不需要大整数范围的场景下，建议使用INT32
   • 可以通过性能测试确定最适合业务需求的数据类型

总结

TensorRT 10.0对INT64数据类型的支持为处理大整数范围的深度学习应用提供了更多可能性。Triton推理服务器通过版本更新及时集成了这一特性，使开发者能够更灵活地部署各类模型。在实际应用中，开发者应当根据具体需求选择合适的数据类型，并确保框架版本、模型配置和输入数据三者之间的严格一致，这是保证推理服务稳定运行的关键所在。