AOTInductor：Torch.Export导出模型的 ahead-of-time 编译

警告

AOTInductor及其相关功能处于原型阶段，可能会破坏向后兼容性。

AOTInductor 是 TorchInductor 的一个专用版本，用于处理导出的 PyTorch 模型、优化它们，并生成共享库和其他相关工件。这些编译后的工件专门设计用于在非 Python 环境中部署，通常应用于服务器端的推理任务。

在这个教程中，你将学习如何将一个 PyTorch 模型导出，编译成共享库，并用 C++ 进行模型预测。

模型编译

使用AOTInductor，你可以用Python编写模型。以下示例展示了如何调用aot_compile将模型转换为共享库。

此 API 使用 torch.export 将模型捕获到计算图中，然后使用 TorchInductor 生成一个可以在非 Python 环境下运行的 .so 文件。有关 torch._export.aot_compile API 的详细信息，请参阅代码此处。有关 torch.export 的更多详细信息，请参阅torch.export 文档。

注意

如果你的机器上有支持CUDA的设备，并且你在安装PyTorch时启用了CUDA支持，以下代码将把模型编译为一个共享库以在CUDA上执行。否则，编译后的程序将在CPU上运行。为了在CPU推理期间获得更好的性能，建议通过设置export TORCHINDUCTOR_FREEZING=1来启用冻结功能，在运行下面的Python脚本之前进行此设置。

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)
    so_path = torch._export.aot_compile(
        model,
        example_inputs,
        # Specify the first dimension of the input x as dynamic
        dynamic_shapes={"x": {0: batch_dim}},
        # Specify the generated shared library path
        options={"aot_inductor.output_path": os.path.join(os.getcwd(), "model.so")},
    )

在这个示例中，Dim 参数用于指定输入变量“x”的第一个维度为动态。需要注意的是，编译库的路径和名称没有被指定，因此共享库会被存储在临时目录中。为了从 C++ 代码中访问这个路径，我们将它保存到一个文件中，并在后续的 C++ 代码中进行检索。

C++中的推断

接下来，我们将使用 inference.cpp 这个 C++ 文件来加载上一步生成的共享库，在 C++ 环境中直接进行模型预测。

注意

以下代码片段假设你的系统有一个支持 CUDA 的设备，并且模型已编译为在 CUDA 上运行（如前所示）。如果没有 GPU，则需要进行如下更改以使其在 CPU 上运行：1. 将 model_container_runner_cuda.h 更改为 model_container_runner_cpu.h 2. 将 AOTIModelContainerRunnerCuda 更改为 AOTIModelContainerRunnerCpu 3. 将 at::kCUDA 更改为 at::kCPU

#include<iostream>
#include<vector>

#include<torch/torch.h>
#include<torch/csrc/inductor/aoti_runner/model_container_runner_cuda.h>

intmain(){
c10::InferenceModemode;

torch::inductor::AOTIModelContainerRunnerCudarunner("model.so");
std::vector<torch::Tensor>inputs={torch::randn({8,10},at::kCUDA)};
std::vector<torch::Tensor>outputs=runner.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.so.
std::cout<<"Result from the second inference:"<<std::endl;
std::vector<torch::Tensor>inputs2={torch::randn({2,10},at::kCUDA)};
std::cout<<runner.run(inputs2)[0]<<std::endl;

return0;
}

为了构建 C++ 文件，你可以使用提供的 CMakeLists.txt 文件。这个文件自动化了调用 python model.py 来进行 AOT 编译模型，并将 inference.cpp 编译成一个名为 aoti_example 的可执行二进制文件。

cmake_minimum_required(VERSION3.18FATAL_ERROR)
project(aoti_example)

find_package(TorchREQUIRED)

add_executable(aoti_exampleinference.cppmodel.so)

add_custom_command(
OUTPUTmodel.so
COMMANDpython${CMAKE_CURRENT_SOURCE_DIR}/model.py
DEPENDSmodel.py
)

target_link_libraries(aoti_example"${TORCH_LIBRARIES}")
set_property(TARGETaoti_examplePROPERTYCXX_STANDARD17)

假设目录结构如下所示，你可以执行以下命令来构建二进制文件。需要注意的是，CMAKE_PREFIX_PATH 变量对于 CMake 查找 LibTorch 库非常重要，并且应该设置为绝对路径。请记住，你的路径可能与示例中的不同。

aoti_example/
CMakeLists.txt
inference.cpp
model.py

$mkdirbuild
$cdbuild
$CMAKE_PREFIX_PATH=/path/to/python/install/site-packages/torch/share/cmakecmake..
$cmake--build.--configRelease

当aoti_example二进制文件生成在build目录后，运行它会显示出类似下面的结果：

$./aoti_example
Resultfromthefirstinference:
0.4866
0.5184
0.4462
0.4611
0.4744
0.4811
0.4938
0.4193
[CUDAFloatType{8,1}]
Resultfromthesecondinference:
0.4883
0.4703
[CUDAFloatType{2,1}]