Shortcuts

torch.compiler

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

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

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

  • TorchDynamo (torch._dynamo) 是一个内部API,它使用称为帧评估API的CPython特性来安全地捕获PyTorch图。PyTorch用户可用的方法通过torch.compiler命名空间公开。

  • TorchInductor 是默认的 torch.compile 深度学习编译器, 能够为多个加速器和后端生成快速代码。您需要使用后端编译器,以通过 torch.compile 实现加速。对于 NVIDIA 和 AMD 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 进行推理优化。阅读更多