torch.compiler#

创建日期：2023 年 7 月 28 日 | 最后更新日期：2025 年 6 月 6 日

torch.compiler 是一个命名空间，通过它，一些内部编译器方法被公开供用户使用。这个命名空间中的主要函数和特性是 torch.compile。

torch.compile 是 PyTorch 2.x 中引入的一个 PyTorch 函数，旨在解决 PyTorch 中精确图捕获的问题，并最终使软件工程师能够更快地运行他们的 PyTorch 程序。torch.compile 用 Python 编写，标志着 PyTorch 从 C++ 向 Python 的过渡。

torch.compile 利用了以下底层技术：

TorchDynamo (torch._dynamo) 是一个内部 API，它使用 CPython 的帧评估 API 来安全地捕获 PyTorch 图。PyTorch 用户可从外部访问的方法通过 torch.compiler 命名空间公开。
TorchInductor 是 torch.compile 的默认深度学习编译器，它为多种加速器和后端生成快速代码。您需要使用后端编译器才能通过 torch.compile 实现加速。对于 NVIDIA、AMD 和 Intel GPU，它利用 OpenAI Triton 作为关键构建模块。
AOT Autograd 不仅捕获用户级代码，还捕获反向传播，从而实现“提前”捕获反向传播。这使得使用 TorchInductor 加速正向和反向传播成为可能。

注意

在某些情况下，本文档中的 torch.compile、TorchDynamo、torch.compiler 可能互换使用。

如上所述，为了更快地运行您的工作流，torch.compile 通过 TorchDynamo 需要一个后端，将捕获的图转换为快速机器代码。不同的后端可以带来不同的优化增益。默认后端称为 TorchInductor，也称为 *inductor*，TorchDynamo 有一个由我们的合作伙伴开发的受支持后端列表，可以通过运行 torch.compiler.list_backends() 来查看，每个后端都有其可选依赖项。

一些最常用的后端包括：

训练与推理后端

后端	描述
`torch.compile(m, backend="inductor")`	使用 TorchInductor 后端。阅读更多
`torch.compile(m, backend="cudagraphs")`	带有 AOT Autograd 的 CUDA 图。阅读更多
`torch.compile(m, backend="ipex")`	在 CPU 上使用 IPEX。阅读更多
`torch.compile(m, backend="onnxrt")`	使用 ONNX Runtime 在 CPU/GPU 上进行训练。阅读更多

仅推理后端

后端	描述
`torch.compile(m, backend="tensorrt")`	使用 Torch-TensorRT 进行推理优化。需要在调用脚本中 `import torch_tensorrt` 以注册后端。阅读更多
`torch.compile(m, backend="ipex")`	在 CPU 上使用 IPEX 进行推理。阅读更多
`torch.compile(m, backend="tvm")`	使用 Apache TVM 进行推理优化。阅读更多
`torch.compile(m, backend="openvino")`	使用 OpenVINO 进行推理优化。阅读更多

torch.compiler#

阅读更多#

文档

教程

资源