这篇文章主要介绍了在使用 TensorRT 编译基于 Vision Transformer（ViT）的模型时遇到的问题及解决过程，包含以下主要内容：

• 问题背景：在工作中使用 TensorRT 编译 ML 模型时，切换到新的主版本时遇到意外问题，测试中视频输入模型后输出完全错误，虽然有测试捕捉到了回归问题，但仍不知如何解决，于是开始了不寻常的调试之旅。

• 模型介绍：

  • Vision Transformer：现代图像和视频模型的常见架构，接受输入视频（作为位图图像数组）并输出分类。文中的测试失败模型基于VideoMAE V2 论文和代码。
  • TensorRT：NVIDIA 提供的高性能深度学习工具和库集合，用于构建 TensorRT“引擎”（编译模型准备加载到 GPU 进行推理），常见步骤为将模型导出为ONNX文件，然后使用trtexec构建优化引擎。

• 调试过程：

  • 设置示例：搭建一个最小化（且失败）的示例，包括获取视频张量、调整视频尺寸、创建模型并加载权重、定义新模型等步骤，最终成功进行了视频分类预测。
  • 使用 ONNX 和 TensorRT 加速推理：将模型导出为 ONNX 文件，然后使用torch-tensorrt构建 TensorRT 引擎，速度提升约 7 倍，但结果出现错误。
  • 二分查找调试：采用二分查找的方法，沿着模型架构逐步 bisect，找出问题所在。首先在 Transformer 编码器和分类头之间进行 bisect，发现问题出现在 Transformer 块中；接着在 Transformer 块内部的 Attention 层和 MLP 层之间进行 bisect，确定问题在 Attention 层。
  • 问题原因及解决：原 VideoMAE 代码未使用F.scaled_dot_product_attention函数，而是使用了来自“Attention Is All You Need”论文的显式公式，在清理代码时引入了F.scaled_dot_product_attention，导致 TensorRT 编译器出现问题。通过将 Attention 层替换为nn.MultiheadAttention并添加向后兼容性代码，恢复了性能并生成了能正常工作的引擎。

• 附录内容：

  • 性能测量：测量不同注意力版本、精度配置和平台架构下的模型运行时性能，包括 Explicit、Scaled-dot-product、Functional MHA 和 Layered MHA 等版本，以及使用torch.set_float32_matmul_precision和model = model.half()等优化方法的效果。
  • 批处理测试：测量不同批量大小下的 TensorRT 推理吞吐量，发现批量大小增加时吞吐量会下降。
  • 模型注意事项：指出文中简化的HelloViT模型存在问题，如视频稀释导致位置嵌入不匹配，以及对于长视频处理的考虑等。

总之，通过二分查找和代码修改，成功解决了使用 TensorRT 编译 ViT 模型时的问题，同时也展示了调试 ML 模型的方法和注意事项。