• 文档 >
  • torch.compiler
快捷键

torch.compiler ¬

torch.compiler 是一个命名空间,通过该命名空间,一些内部编译器方法被暴露给用户使用。该命名空间的主要功能是 torch.compile

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

利用以下底层技术:

  • TorchDynamo(torch._dynamo)是一个内部 API,它使用 CPython 的一个功能——帧评估 API,以安全地捕获 PyTorch 图。通过 torch.compiler 命名空间公开了外部可用的 PyTorch 用户方法。

  • 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")

使用 IPEX 在 CPU 上进行推理。了解更多

torch.compile(m, backend="tvm")

使用 Apache TVM 进行推理优化。了解更多

torch.compile(m, backend="openvino")

使用 OpenVINO 进行推理优化。了解更多

阅读更多 ¶

PyTorch 后端供应商指南


© 版权所有 PyTorch 贡献者。

使用 Sphinx 构建,主题由 Read the Docs 提供。

文档

查看 PyTorch 的全面开发者文档

查看文档

教程

深入了解初学者和高级开发者的教程

查看教程

资源

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

查看资源