DEV Community
·
torch.compile 与 torch.export
内容提要
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.compile()作为JIT编译器，能够灵活处理不可追踪的代码部分。这意味着在开发过程中，用户可以快速迭代和调试，而不必担心编译过程中的限制。这种灵活性适合需要频繁修改模型的研究和实验阶段。
torch.export的部署优势
torch.export作为AOT编译器，能够捕获完整的计算图并生成可移植的ATen级别图形，这使得它在模型部署时具有明显优势。对于需要在生产环境中稳定运行的应用，torch.export提供了更高的可靠性和性能保障。
选择工具的考虑因素
在选择使用torch.compile()还是torch.export时，开发者需要考虑项目的需求。如果项目需要快速迭代和灵活性，torch.compile()可能更合适；而对于需要高效部署和稳定性的项目，torch.export则是更好的选择。
延伸问答
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则捕获完整图形，并在遇到不可追踪代码时报错。
