AOTInductor：Torch.Export 模型的提前编译¶
警告
AOTInductor 及其相关功能处于原型阶段，可能会发生破坏向后兼容性的更改。
AOTInductor 是 TorchInductor 的一个专门版本，旨在处理导出的 PyTorch 模型，对其进行优化，并生成共享库以及其他相关工件。这些编译后的工件专门用于在非 Python 环境中部署，这些环境通常用于服务器端的推理部署。
在本教程中，您将深入了解如何获取 PyTorch 模型、导出它、将其编译成共享库以及使用 C++ 进行模型预测的过程。
模型编译¶
使用 AOTInductor，您仍然可以在 Python 中编写模型。以下示例演示了如何调用 aot_compile 将模型转换为共享库。
此 API 使用 torch.export 将模型捕获到计算图中，然后使用 TorchInductor 生成一个 .so 文件，该文件可以在非 Python 环境中运行。有关 torch._export.aot_compile API 的完整详细信息，您可以参考此处的代码 此处。有关 torch.export 的更多详细信息，您可以参考 torch.export 文档。
注意
如果您的机器上启用了 CUDA 的设备，并且您使用 CUDA 支持安装了 PyTorch，则以下代码将把模型编译成一个用于 CUDA 执行的共享库。否则，编译后的工件将在 CPU 上运行。为了在 CPU 推理期间获得更好的性能，建议在运行下面的 Python 脚本之前启用冻结，方法是设置 export TORCHINDUCTOR_FREEZING=1。
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++ 中进行推理¶
接下来，我们使用以下 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
#include <iostream>
#include <vector>
#include <torch/torch.h>
#include <torch/csrc/inductor/aoti_runner/model_container_runner_cuda.h>
int main() {
c10::InferenceMode mode;
torch::inductor::AOTIModelContainerRunnerCuda runner("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;
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
在 aoti_example 二进制文件生成在 build 目录中后，执行它将显示类似于以下结果
$ ./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} ]
