Transformers.js项目中flan-t5-small模型ONNX转换与推理问题解析

在自然语言处理领域，Google发布的flan-t5-small模型因其轻量级和高效性受到广泛关注。本文将深入分析该模型在使用Transformers.js工具链进行ONNX格式转换时遇到的一个典型问题，以及解决方案。

问题现象

当开发者尝试将flan-t5-small模型转换为ONNX格式并在TensorFlow.js环境中运行时，观察到模型输出出现异常行为：模型在生成第一个token后会不断重复该词，无法继续生成有意义的后续内容。值得注意的是，同样的模型在Python环境中使用ONNXRuntime运行时表现正常。

技术背景

T5（Text-to-Text Transfer Transformer）模型采用编码器-解码器架构，在序列生成任务中需要处理自回归特性。在ONNX转换过程中，关键点在于正确处理模型的缓存机制（past key values），这对生成质量至关重要。

问题根源

经过分析，问题出在模型转换时的任务指定方式。开发者最初使用了：

python -m convert --model_id flan-t5-small --task text2text-generation

这种指定方式生成的ONNX模型不包含past key values输入节点，导致在TensorFlow.js环境中无法正确处理自回归生成过程。

解决方案

正确的转换命令应为：

python -m convert --model_id google/flan-t5-small --task text2text-generation-with-past

或者简化为：

python -m convert --model_id google/flan-t5-small

关键区别在于：

1. 必须使用完整模型ID（包含google/前缀）
2. 明确指定支持past key values的生成任务

技术要点

1. Past Key Values机制：在自回归生成中，模型需要缓存先前时间步的key和value以提升效率。ONNX模型需要显式支持这一机制。

2. 任务类型选择：

   • text2text-generation：基础生成任务，不支持缓存
   • text2text-generation-with-past：支持缓存机制的生成任务

3. 模型ID规范：使用完整模型路径（如google/flan-t5-small）可确保任务自动推断功能正常工作。

最佳实践建议

1. 对于序列生成任务，始终优先使用支持past key values的转换方式
2. 在本地转换时，建议显式指定任务类型
3. 转换完成后，应验证模型在目标环境中的生成质量
4. 对于生产环境，考虑使用量化技术减小模型体积

通过正确理解和使用模型转换参数，开发者可以确保flan-t5-small等序列生成模型在各种运行环境中保持一致的生成质量。这一经验同样适用于其他类似架构的模型转换工作。