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 GPU,它利用 OpenAI Triton 作为关键构建模块。AOT Autograd 不仅捕获用户级代码,还捕获反向传播,从而导致“提前”捕获反向传递。这使得能够使用 TorchInductor 加速正向和反向传递。
注意
在某些情况下,术语 torch.compile
、TorchDynamo、torch.compiler
可能在此文档中互换使用。
如上所述,为了更快地运行工作流,torch.compile
通过 TorchDynamo 需要一个后端,该后端将捕获的图形转换为快速的机器代码。不同的后端可以产生不同的优化收益。默认后端称为 TorchInductor,也称为inductor,TorchDynamo 有一系列由我们的合作伙伴开发的支持后端,可以通过运行 torch.compiler.list_backends()
查看,每个后端都有其可选依赖项。
一些最常用的后端包括
训练和推理后端
后端 |
说明 |
---|---|
|
使用 TorchInductor 后端。 了解更多 |
|
具有 AOT Autograd 的 CUDA 图形。 了解更多 |
|
在 CPU 上使用 IPEX。 了解更多 |
|
在 CPU/GPU 上使用 ONNX Runtime 进行训练。 了解更多 |
仅推理后端
后端 |
说明 |
---|---|
|
使用 Torch-TensorRT 进行推理优化。需要在调用脚本中 |
|
在 CPU 上使用 IPEX 进行推理。 了解更多 |
|
使用 Apache TVM 进行推理优化。 了解更多 |
|
使用 OpenVINO 进行推理优化。 了解更多 |