torch.compiler¶

torch.compiler 是一个命名空间，通过该命名空间可以访问一些内部编译器方法。此命名空间的主要功能和特性是 torch.compile。

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

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

TorchDynamo (torch._dynamo) 是一个内部API，它使用CPython的一个称为Frame Evaluation API的功能来安全地捕获PyTorch图。可供PyTorch用户使用的外部方法通过torch.compiler命名空间提供。
TorchInductor 是默认的 torch.compile 深度学习编译器，它为多个加速器和后端生成快速代码。你需要使用后端编译器来通过 torch.compile 实现加速。对于 NVIDIA、AMD 和 Intel GPU，它利用 OpenAI Triton 作为关键构建块。
AOT 自动微分 不仅捕获了用户级别的代码，还捕获了反向传播，这导致在“提前”阶段捕获反向传递。这使得可以使用 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")`	CUDA 图与 AOT Autograd。了解更多
`torch.compile(m, backend="ipex")`	在CPU上使用IPEX。了解更多
`torch.compile(m, backend="onnxrt")`	使用ONNX运行时在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进行推理优化。了解更多

了解更多¶

PyTorch 用户入门指南

深入探索PyTorch开发者

PyTorch后端供应商入门指南

torch.compiler¶

了解更多¶

文档

教程

资源