评价此页
★ ★ ★ ★ ★

AOTInductor：针对 Torch.Export 模型进行提前（Ahead-Of-Time）编译#

创建日期：2025 年 6 月 13 日 | 最后更新日期：2025 年 12 月 3 日

AOTInductor 是 TorchInductor 的专用版本，旨在处理导出的 PyTorch 模型，对其进行优化，并生成共享库及其他相关制品。这些编译后的制品专为非 Python 环境下的部署而设计，常用于服务器端的推理部署。

在本教程中，您将深入了解如何获取 PyTorch 模型、将其导出、编译为制品，并使用 C++ 执行模型预测的过程。

模型编译#

要使用 AOTInductor 编译模型，我们首先需要使用 torch.export.export() 将给定的 PyTorch 模型捕获为计算图。torch.export 为捕获的 IR（中间表示）提供了稳健性保证和严格的规范，AOTInductor 正是依赖于此。

然后，我们将使用 torch._inductor.aoti_compile_and_package() 通过 TorchInductor 编译导出的程序，并将编译后的制品保存到一个包中。该包采用 PT2 归档规范 (PT2 Archive Spec) 格式。

注意

如果您的机器上有支持 CUDA 的设备，并且安装了支持 CUDA 的 PyTorch，则以下代码会将模型编译为用于 CUDA 执行的共享库。否则，编译后的制品将在 CPU 上运行。为了在 CPU 推理期间获得更好的性能，建议在运行以下 Python 脚本之前通过设置 export TORCHINDUCTOR_FREEZING=1 来启用冻结（freezing）。此行为同样适用于 Intel® GPU 环境。

import os
import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(10, 16)
        self.relu = torch.nn.ReLU()
        self.fc2 = torch.nn.Linear(16, 1)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return x

with torch.no_grad():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = Model().to(device=device)
    example_inputs=(torch.randn(8, 10, device=device),)
    batch_dim = torch.export.Dim("batch", min=1, max=1024)
    # [Optional] Specify the first dimension of the input x as dynamic.
    exported = torch.export.export(model, example_inputs, dynamic_shapes={"x": {0: batch_dim}})
    # [Note] In this example we directly feed the exported module to aoti_compile_and_package.
    # Depending on your use case, e.g. if your training platform and inference platform
    # are different, you may choose to save the exported model using torch.export.save and
    # then load it back using torch.export.load on your inference platform to run AOT compilation.
    output_path = torch._inductor.aoti_compile_and_package(
        exported,
        # [Optional] Specify the generated shared library path. If not specified,
        # the generated artifact is stored in your system temp directory.
        package_path=os.path.join(os.getcwd(), "model.pt2"),
        # [Optional] Specify Inductor configs
        # This specific max_autotune option will turn on more extensive kernel autotuning for
        # better performance.
        inductor_configs={"max_autotune": True,},
    )

在这个说明性示例中，Dim 参数被用来将输入变量“x”的第一维指定为动态维度。值得注意的是，编译库的路径和名称未指定，因此共享库将存储在临时目录中。为了从 C++ 端访问此路径，我们将其保存到文件中，以便稍后在 C++ 代码中检索。

在 Python 中进行推理#

部署编译后的制品以进行推理有多种方式，其中之一是使用 Python。我们在 Python 中提供了一个便捷的实用 API torch._inductor.aoti_load_package() 用于加载和运行该制品，如下例所示：

import os
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
model = torch._inductor.aoti_load_package(os.path.join(os.getcwd(), "model.pt2"))
print(model(torch.randn(8, 10, device=device)))

推理时的输入应与导出时的输入具有相同的大小、数据类型和步长（stride）。

在 C++ 中进行推理#

接下来，我们使用以下示例 C++ 文件 inference.cpp 来加载编译后的制品，使我们能够直接在 C++ 环境中进行模型预测。

#include <iostream>
#include <vector>

#include <torch/torch.h>
#include <torch/csrc/inductor/aoti_package/model_package_loader.h>

int main() {
    c10::InferenceMode mode;

    torch::inductor::AOTIModelPackageLoader loader("model.pt2");
    // Assume running on CUDA
    std::vector<torch::Tensor> inputs = {torch::randn({8, 10}, at::kCUDA)};
    std::vector<torch::Tensor> outputs = loader.run(inputs);
    std::cout << "Result from the first inference:"<< std::endl;
    std::cout << outputs[0] << std::endl;

    // The second inference uses a different batch size and it works because we
    // specified that dimension as dynamic when compiling model.pt2.
    std::cout << "Result from the second inference:"<< std::endl;
    // Assume running on CUDA
    std::cout << loader.run({torch::randn({1, 10}, at::kCUDA)})[0] << std::endl;

    return 0;
}

为了构建 C++ 文件，您可以利用提供的 CMakeLists.txt 文件，它会自动调用 python model.py 对模型进行 AOT 编译，并将 inference.cpp 编译为名为 aoti_example 的可执行二进制文件。

cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
project(aoti_example)

find_package(Torch REQUIRED)

add_executable(aoti_example inference.cpp model.pt2)

add_custom_command(
    OUTPUT model.pt2
    COMMAND python ${CMAKE_CURRENT_SOURCE_DIR}/model.py
    DEPENDS model.py
)

target_link_libraries(aoti_example "${TORCH_LIBRARIES}")
set_property(TARGET aoti_example PROPERTY CXX_STANDARD 17)

如果目录结构类似于以下内容，您可以执行后续命令来构建二进制文件。务必注意，CMAKE_PREFIX_PATH 变量对于 CMake 定位 LibTorch 库至关重要，并且应设置为绝对路径。请注意，您的路径可能与此示例中演示的路径不同。 

aoti_example/
    CMakeLists.txt
    inference.cpp
    model.py

$ mkdir build
$ cd build
$ CMAKE_PREFIX_PATH=/path/to/python/install/site-packages/torch/share/cmake cmake ..
$ cmake --build . --config Release

在 build 目录中生成 aoti_example 二进制文件后，执行它将显示类似于以下的结果：

$ ./aoti_example
Result from the first inference:
0.4866
0.5184
0.4462
0.4611
0.4744
0.4811
0.4938
0.4193
[ CUDAFloatType{8,1} ]
Result from the second inference:
0.4883
0.4703
[ CUDAFloatType{2,1} ]

故障排除#

以下是一些用于调试 AOT Inductor 的实用工具。

调试工具

要启用对输入的运行时检查，请将环境变量 AOTI_RUNTIME_CHECK_INPUTS 设置为 1。如果编译后的模型的输入在大小、数据类型或步长上与导出时使用的不同，这将引发 RuntimeError。

API参考#

torch._inductor.aoti_compile_and_package(exported_program, _deprecated_unused_args=None, _deprecated_unused_kwargs=None, *, package_path=None, inductor_configs=None)[source]#

使用 AOTInductor 编译导出程序，并将其打包成由输入 package_path 指定的 .pt2 制品。要加载该包，您可以调用 torch._inductor.aoti_load_package(package_path)。

用法示例如下：

ep = torch.export.export(M(), ...)
aoti_file = torch._inductor.aoti_compile_and_package(
    ep, package_path="my_package.pt2"
)
compiled_model = torch._inductor.aoti_load_package("my_package.pt2")

要将多个模型编译并保存到单个 .pt2 制品中，您可以执行以下操作：

ep1 = torch.export.export(M1(), ...)
aoti_file1 = torch._inductor.aot_compile(
    ep1, ..., options={"aot_inductor.package": True}
)
ep2 = torch.export.export(M2(), ...)
aoti_file2 = torch._inductor.aot_compile(
    ep2, ..., options={"aot_inductor.package": True}
)

from torch._inductor.package import package_aoti, load_package

package_aoti("my_package.pt2", {"model1": aoti_file1, "model2": aoti_file2})

compiled_model1 = load_package("my_package.pt2", "model1")
compiled_model2 = load_package("my_package.pt2", "model2")
参数:

  • exported_program (ExportedProgram) – 通过调用 torch.export 创建的导出程序。

  • package_path (Optional[FileLike]) – 生成的 .pt2 制品的可选指定路径。

  • inductor_configs (Optional[dict[str, Any]]) – 用于控制 inductor 的可选配置字典。

返回:

生成的制品的路径。

返回类型:

str

torch._inductor.aoti_load_package(path, run_single_threaded=False, device_index=-1)[source]#

从 PT2 包中加载模型。

如果 PT2 中打包了多个模型，这将加载默认模型。要加载特定模型，您可以直接调用加载 API。 

from torch._inductor.package import load_package

compiled_model1 = load_package("my_package.pt2", "model1")
compiled_model2 = load_package("my_package.pt2", "model2")
参数:

  • path (FileLike) – .pt2 包的路径。

  • run_single_threaded (bool) – 是否应在没有线程同步逻辑的情况下运行模型。这对于避免与 CUDAGraphs 的冲突非常有用。

  • device_index (int) – 要加载 PT2 包的设备索引。默认使用 device_index=-1，这在使用 CUDA 时对应于 cuda 设备。例如，传递 device_index=1 会将包加载到 cuda:1。

返回类型:

AOTICompiledModel

在此页面上
在 GitHub 上编辑
显示源代码
PyTorch 库