PyTorch AOTInductor 实战Torch.Export 模型的提前编译、打包与 C 部署完整指南【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchAOTInductorAhead-Of-Time Inductor是 PyTorch 中 TorchInductor 的提前编译版本它接收通过torch.export.export导出的计算图经优化后产出生成本地共享库shared library及配套产物最终以.pt2归档格式打包专为非 Python 环境下的服务端推理部署而设计。本文将基于官方用户指南与当前仓库源码完整演示从模型导出、AOT 编译打包到 Python 侧加载与 C 侧加载推理、CMake 工程构建的端到端流程并深入讲解动态维度声明、Inductor 配置、运行时输入检查等关键细节帮助你掌握一套可在生产环境落地的 PyTorch 部署方案。什么是 AOTInductor从 Define-by-Run 到提前编译AOTInductor 是 TorchInductor 的专门化版本。TorchInductor 是 PyTorch 原生的编译器采用 define-by-run IR 与符号形状symbolic shapes机制而 AOTInductor 的定位是处理已被导出的exportedPyTorch 模型先做优化再产出生成本地共享库及其他相关工件artifacts。这些编译产物专门用于非 Python 环境下的部署最常见的场景就是服务端推理inference deployment。换句话说AOTInductor 让训练/导出发生在 Python 世界、部署运行在 C/无 Python 运行时环境成为可能从而绕开 Python 解释器带来的启动开销与依赖问题。从当前仓库的源码结构看AOTInductor 的完整实现横跨两个层面Python 侧入口torch/_inductor/init.py 中的aoti_compile_and_package、aoti_load_package、aot_compile等 APIC 侧运行时torch/csrc/inductor/aoti_package/model_package_loader.h 中的AOTIModelPackageLoader类以及其底层依赖的AOTIModelContainerRunner见 torch/csrc/inductor/aoti_runner/model_container_runner.h。整条链路可以概括为三步导出用torch.export.export把nn.Module捕获成一张计算图ExportedProgram。torch.export对捕获的 IR 提供 soundness 保证和严格的规范这正是 AOTInductor 能安全编译的前提编译并打包用torch._inductor.aoti_compile_and_package调用 TorchInductor 完成编译并把产物打包成一个.pt2文件遵循 PT2 Archive Spec即 export.pt2_archive 描述的归档规范加载推理Python 侧通过torch._inductor.aoti_load_packageC 侧通过AOTIModelPackageLoader加载.pt2归档并执行推理。模型编译导出、动态形状声明与 AOT 打包第一步用 torch.export.export 捕获计算图编译的第一步是使用torch.export.export把模型捕获为计算图。之所以必须以导出结果为输入是因为torch.export提供的 IR 规范严格且具备 soundness 保证——AOTInductor 的代码生成依赖这一约束因此在aoti_compile_and_package的实现中会首先校验输入类型只有ExportedProgram实例才会被接受否则抛出ValueError(Only ExportedProgram is supported)见 torch/_inductor/init.py。下面的model.py示例构建了一个两层全连接网络并在torch.no_grad()下完成导出与编译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(devicedevice) example_inputs(torch.randn(8, 10, devicedevice),) batch_dim torch.export.Dim(batch, min1, max1024) # [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_pathos.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,}, )关于这段代码有几点值得展开动态维度dynamic shapestorch.export.Dim(batch, min1, max1024)声明了输入x的第 0 维是动态的取值范围为[1, 1024]。这样编译出来的产物在推理时就能接受不同 batch size 的输入下文 C 示例会用 batch8 和 batch1 两次推理来验证这一点。如果不声明动态维度导出与推理时的张量形状必须完全一致导出与编译解耦示例直接把这个仓库中未编译的exported模块喂给aoti_compile_and_package。如果训练平台与推理平台分离可以先在本平台用torch.export.save保存导出结果再到推理平台用torch.export.load加载后执行 AOT 编译产物位置package_path指定生成的.pt2归档路径不指定时归档会写入系统临时目录。为了让 C 侧能拿到路径示例把它写进了文件以便后续读取Inductor 配置inductor_configs{max_autotune: True}会开启更全面的内核自动调优autotuning以换取更好的性能。第二步aoti_compile_and_package 内部发生了什么在仓库源码 torch/_inductor/init.py 中aoti_compile_and_package的实现揭示了几条关键约束校验导出程序输入必须是ExportedProgram且exported_program.example_inputs必须已设置否则抛出RuntimeError提示 AOTInductor 编译需要示例输入校验 package_path路径必须以.pt2结尾或为None走系统临时目录也可以是一个可写、可 seek 的文件缓冲区io.IOBase/IO。否则抛出AssertionError(Expect package path to be a file ending in .pt2, ...)强制开启打包模式函数内部会自动设置inductor_configs[aot_inductor.package] True并拒绝同时指定aot_inductor.output_path配置——打包路径应该通过package_path参数传入minifier 保护整个编译过程被aot_inductor_minifier_wrapper包裹来自 torch/_inductor/debug.py一旦编译失败会自动触发最小化复现minifier流程便于定位问题返回产物路径返回值是生成的.pt2归档路径字符串。此外仓库还提供了更底层的aot_compileAPItorch/_inductor/init.py它直接接受torch.fx.GraphModule通过compile_fx_aot编译为共享库并支持options{aot_inductor.package: True}与package_aoti配合把多个模型打包进同一个.pt2归档详见下文多模型打包。设备与性能提示官方文档特别提示如果机器上有 CUDA 设备且安装的 PyTorch 带 CUDA 支持上述代码会把模型编译为面向 CUDA 执行的共享库否则编译产物在 CPU 上运行。Intel GPU 环境的行为相同为了获得更好的 CPU 推理性能建议在运行上述 Python 脚本前开启冻结freezingexport TORCHINDUCTOR_FREEZING1。冻结的语义在 torch/_inductor/config.py 中有明确注释冻结时会抓取模型参数并执行常量折叠constant folding等优化之后这些权重不再作为图输入从而减少运行时开销freezing_discard_parameters默认False用于控制是否丢弃 eager 模式下的nn.Module参数以降低内存开销。TORCHINDUCTOR_MAX_AUTOTUNE1环境变量则可以等价地开启max_autotune见 config.py。Python 侧推理aoti_load_package编译产物有多种部署方式其中一种是直接用 Python 加载推理。仓库为此提供了便捷工具 APItorch._inductor.aoti_load_packageimport 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, devicedevice)))aoti_load_package的完整签名torch/_inductor/init.py为aoti_load_package(path, run_single_threadedFalse, device_index-1) - AOTICompiledModelpath.pt2归档路径或已解包的归档目录路径run_single_threaded是否在无线程同步逻辑下运行便于与 CUDAGraph 等场景共存避免冲突device_index产物加载到的设备索引。默认-1表示 CUDA 环境下的cuda传1则加载到cuda:1。从源码可以看到torch/_inductor/package/package.pyload_package内部对路径做了区分如果传的是目录则视为已解包目录直接构造 loader如果传的是.pt2归档则先解包再通过torch._C._aoti.AOTIModelPackageLoaderC 绑定的 loader创建AOTICompiledModel包装对象。若读取失败且不是设备相关错误会打印 Loading outdated pt2 file. Please regenerate your package. 警告并回退到兼容路径。注意推理时的输入必须与导出时保持相同的 size、dtype 与 stride。这是 AOT 编译的固有约束——代码生成基于导出时的张量元数据完成。C 侧推理AOTIModelPackageLoader 与动态 batch接下来是本文的核心场景——在纯 C 环境中加载编译产物做推理。官方指南提供了inference.cpp#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::vectortorch::Tensor inputs {torch::randn({8, 10}, at::kCUDA)}; std::vectortorch::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; }要点说明c10::InferenceMode mode;进入推理模式禁用梯度跟踪减少运行时开销AOTIModelPackageLoader构造函数默认加载名为model的模型见 model_package_loader.h其完整签名还支持model_name、run_single_threaded、num_runnersrunner 数量与device_index参数loader.run(inputs)接收std::vectorat::Tensor输入并返回输出张量向量底层还提供boxed_run会偷走输入张量所有权、load_constants加载常量到模型设备allow_h2d_copyTrue时允许 CPU 常量静默拷贝到模型设备、get_metadata/get_call_spec/get_constant_fqns等用于查询与更新常量的方法第二次推理换用{1, 10}的 batch 依旧成功正是因为编译时通过torch.export.Dim(batch, min1, max1024)把该维声明为动态。构建 C 可执行文件CMake 工程官方提供了配套的CMakeLists.txt它自动化完成两件事调用python model.py执行 AOT 编译生成model.pt2以及把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})目录结构约定如下aoti_example/ CMakeLists.txt inference.cpp model.py构建命令如下。关键点CMAKE_PREFIX_PATH必须设置为绝对路径它用于让 CMake 定位 LibTorch 库示例中的路径仅为示意实际路径可能不同通常是 Python 安装目录下的site-packages/torch/share/cmake$ 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} ]注意示例输出中第二次推理显示的 shape 为{2,1}这只是演示脚本中随机张量在文档撰写时的快照实际运行会以你传入的 batch size 为准例如{1,10}输入对应{1,1}输出。Troubleshooting调试工具与运行时输入检查官方指南列出了一些实用的 AOTInductor 调试手段对应文档还包括三份独立资料日志配置说明docs/source/logging.rst最小化复现工具docs/source/user_guide/torch_compiler/torch.compiler_aot_inductor_minifier.md调试指南docs/source/user_guide/torch_compiler/torch.compiler_aot_inductor_debugging_guide.md。此外运行时输入检查是一个非常实用的开关设置环境变量AOTI_RUNTIME_CHECK_INPUTS1后如果编译模型的输入在 size、dtype 或 stride 上与导出时不一致会抛出RuntimeError。其底层实现位于 torch/_inductor/codegen/cpp_wrapper_cpu.py代码生成器会在生成的 AOT 源码中注入_check_aoti_runtime_check_inputs_env()辅助函数——只有当环境变量存在且首字符不为0时才执行__check_inputs_outputs对每个图输入调用check_input_{idx}(input_handles)做逐项校验从而在不检查时保持零额外开销。在 config.py 附近还可以看到该机制与动态形状下界如[2, ...]的配合说明避免因放宽的动态形状边界触发误报。进阶多模型打包与产物组织aoti_compile_and_package的 docstring 与 torch/_inductor/package/package.py 展示了多模型合并进单个.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)package_aoti(archive_file, aoti_files)会把若干 AOTInductor 产物按{模型名: 产物路径}字典组织并写入 PT2Archive 格式加载时load_package(path, model_name)按名取用aoti_load_package则默认取名为model的那个。这为多模型服务、AB 测试或模型热更新场景提供了统一的产物管理方式。API 参考本文涉及的官方 API 定义如下完整签名与 docstring 请以源码为准torch._inductor.aoti_compile_and_package(exported_program, *, package_pathNone, inductor_configsNone) - str编译导出程序并以.pt2归档形式打包返回产物路径见 torch/_inductor/init.pytorch._inductor.aoti_load_package(path, run_single_threadedFalse, device_index-1) - AOTICompiledModel从.pt2归档或解包目录加载已编译模型并可直接调用推理见 torch/_inductor/init.py。小结AOTInductor 为 PyTorch 模型提供了一条从Python 导出到本地共享库 PT2 归档再到Python/C 双端加载推理的完整提前编译部署链路。本文覆盖了torch.export.export与torch.export.Dim动态维度声明、aoti_compile_and_package的打包语义与参数约束、Python 侧aoti_load_package的加载方式、C 侧AOTIModelPackageLoader的动态 batch 推理、CMake 工程构建要点、AOTI_RUNTIME_CHECK_INPUTS运行时检查机制以及多模型合并打包的进阶用法。你可以以此为模板把任意导出的模型编译成不依赖 Python 运行时的推理产物直接嵌入 C 服务进程。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考