Sherpa-onnx项目中Whisper模型ONNX转换的类型错误解决方案

在将Whisper语音识别模型转换为ONNX格式时，开发者可能会遇到一个常见的类型错误："TypeError: scaled_dot_product_attention(): argument 'is_causal' must be bool, not Tensor"。这个问题主要出现在PyTorch的注意力机制实现与ONNX导出过程的兼容性上。

问题背景

Whisper模型使用了PyTorch的scaled_dot_product_attention函数来实现注意力机制。在导出为ONNX格式时，该函数的is_causal参数需要严格的布尔类型，但实际传递的却是一个张量(Tensor)类型，导致类型不匹配错误。

根本原因分析

这个问题的根源在于Whisper模型实现中的n_ctx变量。在某些情况下，特别是在动态计算上下文长度时，n_ctx可能被推断为张量类型而非简单的整数类型。当这个值被用于设置注意力机制的因果掩码(is_causal)时，就会导致类型不匹配。

解决方案

方案一：强制类型转换

最直接的解决方案是在Whisper模型的qkv_attention函数中，将n_ctx显式转换为整数类型：

n_ctx = int(n_ctx)  # 确保n_ctx是整数而非张量

这种方法简单有效，但需要修改模型源代码，可能不适合所有使用场景。

方案二：使用最新PyTorch版本

如果使用较新版本的PyTorch(如2.6.0及以上)，这个问题可能已经被修复。PyTorch团队不断改进ONNX导出功能，新版本通常能更好地处理这类类型转换问题。

方案三：禁用SDPA功能

Sherpa-onnx项目提供了一个更优雅的解决方案，不需要修改模型源代码：

from whisper.model import disable_sdpa

if __name__ == '__main__':
    with disable_sdpa():
        main()

这种方法通过上下文管理器临时禁用scaled_dot_product_attention(SDPA)功能，使用传统的注意力实现来避免类型问题。

技术建议

对于需要将Whisper模型导出为ONNX格式的开发者，建议：

1. 首先尝试使用最新稳定版的PyTorch
2. 如果问题仍然存在，优先采用disable_sdpa方案，它最为简洁且不需要修改源代码
3. 只有在特殊情况下才考虑修改模型源代码的方案

这些解决方案不仅适用于Sherpa-onnx项目，对于其他需要将Whisper模型导出为ONNX格式的应用场景也同样有效。