AOTInductor: 针对Torch.Export导出模型的预编译¶

警告

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上运行。

```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,
        # 指定输入x的第一个维度为动态
        dynamic_shapes={"x": {0: batch_dim}},
        # 指定生成的共享库路径
        options={"aot_inductor.output_path": os.path.join(os.getcwd(), "model.so")},
    )
```

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

C++中的推理¶

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

注意

以下代码片段假设您的系统具有支持CUDA的设备，并且您的模型如前所示已编译为在CUDA上运行。在没有GPU的情况下，需要进行这些调整以便在CPU上运行： 1. 将 model_container_runner_cuda.h 更改为 model_container_runner_cpu.h 2. 将 AOTIModelContainerRunnerCuda 更改为 AOTIModelContainerRunnerCpu 3. 将 at::kCUDA 更改为 at::kCPU

```cpp
#include 
#include 

#include 
#include 

int main() {
    c10::InferenceMode mode;

    torch::inductor::AOTIModelContainerRunnerCuda runner("model.so");
    std::vector inputs = {torch::randn({8, 10}, at::kCUDA)};
    std::vector outputs = runner.run(inputs);
    std::cout << "第一次推理的结果:" << std::endl;
    std::cout << outputs[0] << std::endl;

    // 第二次推理使用不同的批次大小，并且它能够正常工作，因为我们
    // 在编译model.so时指定了该维度为动态的。
    std::cout << "第二次推理的结果:" << std::endl;
    std::cout << runner.run({torch::randn({2, 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.so)

add_custom_command(
    OUTPUT model.so
    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
来自第一次推理的结果:
0.4866
0.5184
0.4462
0.4611
0.4744
0.4811
0.4938
0.4193
[ CUDAFloatType{8,1} ]
来自第二次推理的结果:
0.4883
0.4703
[ CUDAFloatType{2,1} ]