DJL项目中TorchScript模型推理性能问题分析

问题背景

在深度学习模型部署过程中，PyTorch的TorchScript格式常被用于提升模型推理性能。然而，在DJL项目使用过程中，开发者发现将HuggingFace的BGE-reranker-v2-m3模型转换为TorchScript格式后，推理速度反而比原始PyTorch模型慢了4-5倍。

性能对比测试

通过对比测试发现，在CPU环境下：

• 原始CrossEncoder模型的10次推理耗时约0.45秒
• TorchScript格式模型的10次推理耗时约2.07秒

进一步测试表明，这种性能差异不仅存在于特定封装模型(CrossEncoder)，在直接使用HuggingFace的AutoModelForSequenceClassification时也同样存在。

CUDA环境下的问题

当尝试在CUDA环境下进行测试时，遇到了更严重的问题。使用PyTorch 2.x版本进行TorchScript转换时，会出现NVRTC编译错误，具体表现为浮点数常量解析失败。这个错误源于PyTorch 2.x版本中的一个已知bug，影响了CUDA内核代码的生成。

解决方案建议

对于遇到类似问题的开发者，可以考虑以下解决方案：

1. 版本回退：对于GPU环境，暂时回退到PyTorch 1.13.1版本可以避免TorchScript的CUDA编译问题。

2. 转换格式选择：考虑使用ONNX格式替代TorchScript，ONNX运行时通常能提供更好的性能表现和更广泛的部署支持。

3. 性能优化：在模型转换前，确保进行充分的预热(warmup)和性能分析，以确定性能瓶颈所在。

技术启示

这一案例揭示了深度学习模型部署中的几个重要考量点：

• 不同模型格式的性能表现可能因具体模型结构和运行环境而异
• 框架版本选择对模型部署成功率和性能有重大影响
• 在实际部署前进行充分的性能测试和验证至关重要

开发者应当根据具体应用场景和性能需求，选择最适合的模型格式和转换方案，而不是盲目追求某种特定的优化技术。