DEV Community
·
torch.compile 与 torch.export
💡
原文英文，约200词，阅读约需1分钟。
📝
内容提要
torch.compile()和torch.export()是基于PyTorch 2.6的工具。前者为JIT编译器，灵活处理不可追踪部分；后者为AOT编译器，捕获完整图形，适合部署。两者在图形捕获、处理不可追踪代码和输出方面存在差异。
🎯
关键要点
- torch.compile()和torch.export()是基于PyTorch 2.6的工具，分别用于不同目的。
- torch.compile()是一个JIT编译器，灵活处理不可追踪部分，回退到Python运行时。
- torch.export是一个AOT编译器，捕获完整图形，对不可追踪代码报错，适合部署。
- torch.compile()使用TorchDynamo进行追踪，AOT Autograd进行优化，torch.fx作为图形基础。
- torch.export依赖TorchDynamo进行广泛的字节码追踪，AOT Autograd进行ATen降级，torch.fx用于图形表示。
- torch.compile()的编译方式为JIT，图形捕获为部分，处理不可追踪代码时回退到Python运行时，输出灵活。
- torch.export的编译方式为AOT，图形捕获为完整，处理不可追踪代码时需要重写或提供信息，输出为可移植的ATen级别图形。
❓
延伸问答
torch.compile()和torch.export()有什么区别？
torch.compile()是JIT编译器，处理不可追踪部分时回退到Python运行时；而torch.export是AOT编译器，捕获完整图形，对不可追踪代码报错。
torch.compile()的主要特点是什么？
torch.compile()使用TorchDynamo进行追踪，灵活处理不可追踪部分，编译方式为JIT，图形捕获为部分。
torch.export适合什么场景使用？
torch.export适合部署场景，因为它捕获完整图形并生成可移植的ATen级别图形。
torch.compile()如何处理不可追踪代码？
torch.compile()在处理不可追踪代码时会回退到Python运行时，提供灵活性。
torch.export的图形输出有什么特点？
torch.export的图形输出为可移植的ATen级别图形，适合独立使用。
torch.compile()和torch.export在图形捕获上有什么不同？
torch.compile()进行部分图形捕获，而torch.export则捕获完整图形，并在遇到不可追踪代码时报错。
➡️
