快捷方式

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,它使用一个名为 Frame Evaluation API 的 CPython 功能来安全地捕获 PyTorch 图。可供 PyTorch 用户在外部使用的这些方法通过 torch.compiler 命名空间公开。

  • **TorchInductor** 是默认的 torch.compile 深度学习编译器,它为多个加速器和后端生成快速代码。您需要使用后端编译器才能通过 torch.compile 才能实现加速。对于 NVIDIA、AMD 和 Intel GPU,它利用 OpenAI Triton 作为关键构建块。

  • **AOT Autograd** 不仅捕获用户级代码,还捕获反向传播,这会导致“提前”捕获反向传播过程。这使得可以使用 TorchInductor 加速前向和反向传播过程。

注意

在某些情况下,术语 torch.compile、TorchDynamo、torch.compiler 在本文档中可能可以互换使用。

如上所述,为了更快地运行您的工作流程,通过 TorchDynamo 使用 torch.compile 需要一个后端,该后端将捕获的图转换为快速的机器代码。不同的后端可以带来不同的优化增益。默认后端称为 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 进行推理优化。了解更多

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源