PyTorch 自定义操作全解析：从实现到导出编译，附代码示例！

引言在 PyTorch 模型中使用自定义操作是很常见的。PyTorch 自定义操作可以是用 C 和 CUDA 实现的自定义类和自定义函数并且可在 Python 和 C 推理程序中使用。在这篇博客文章中将分享如何用 C 和 CUDA 实现 PyTorch 自定义操作以及如何在 PyTorch 模型和 AOTInductor 编译的推理程序中使用它们。这里会以一个简单的恒等卷积示例进行说明。PyTorch 自定义函数PyTorch 自定义函数可以用 C 和 CUDA 实现并使用 TORCH_LIBRARY_IMPL 宏进行注册。可以同时提供 CPU 和 CUDA 实现PyTorch 会根据输入张量的设备自动调度到正确的实现。custom_ops.cpp// ---------------------------------------------------------------------------// CPU 实现通过 clone() 进行简单的逐元素复制。// ---------------------------------------------------------------------------torch::Tensor identity_conv_cpu_impl(const torch::Tensor input){ TORCH_CHECK(!input.is_cuda(), identity_conv_cpu_impl: 输入必须是 CPU 张量); return input.clone();}// ---------------------------------------------------------------------------// 主机端调度器。// ---------------------------------------------------------------------------torch::Tensor identity_conv_cuda_impl(const torch::Tensor input){ TORCH_CHECK(input.is_cuda(), identity_conv_cuda_impl: 输入必须是 CUDA 张量); // 输出与输入具有相同的形状、数据类型和步长。 auto output torch::empty_like(input); const int64_t numel input.numel(); if (numel 0) return output; // 将形状和步长上传到设备以便内核可以读取它们。 const int ndim input.dim(); const auto opts torch::TensorOptions().dtype(torch::kInt64).device(input.device()); const auto shape_dev torch::tensor( std::vector (input.sizes().begin(), input.sizes().end()), opts); const auto strides_dev torch::tensor( std::vector (input.strides().begin(), input.strides().end()), opts); constexpr int kThreads 256; const int blocks static_cast ((numel kThreads - 1) / kThreads); AT_DISPATCH_FLOATING_TYPES_AND2( at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), identity_conv_cuda_impl, []() { identity_kernel ( input.data_ptr (), output.data_ptr (), shape_dev.data_ptr (), strides_dev.data_ptr (), ndim, numel); }); C10_CUDA_KERNEL_LAUNCH_CHECK(); return output;}custom_op_registration.cpp// CUDA 内核实现 my_ops::identity_conv_op。TORCH_LIBRARY_IMPL(my_ops, CUDA, m){ m.impl(identity_conv_op, identity_conv_cuda_impl);}// CPU 回退。TORCH_LIBRARY_IMPL(my_ops, CPU, m){ m.impl(identity_conv_op, identity_conv_cpu_impl);}PyTorch 自定义类PyTorch 自定义函数是无状态的不能持有任何参数。如果想实现一个自定义类使其能够持有一些参数并且有一个可以从 Python 调用的 forward() 方法可以使用 torch::CustomClassHolder 在 C 中定义一个自定义类并使用 TORCH_LIBRARY 宏进行注册。custom_class.cpp// ---------------------------------------------------------------------------// IdentityConvClass// // 这是一个使用 torch.classes 注册的自定义类因此它可以嵌入到 torch.nn.Module 中通过 torch.export 导出并使用 AOTInductor 进行编译。// // forward() 方法委托给 CUDA 恒等内核。channels_ 字段用于语义完整性并通过 def_pickle 进行序列化以便该类在导出/导入过程中保持不变。// ---------------------------------------------------------------------------struct IdentityConvClass : torch::CustomClassHolder{ int64_t channels_; explicit IdentityConvClass(int64_t channels) : channels_(channels) {} torch::Tensor forward(const torch::Tensor x) { return x.is_cuda() ? identity_conv_cuda_impl(x) : identity_conv_cpu_impl(x); } int64_t get_channels() const { return channels_; };custom_class_registration.cpp// ---------------------------------------------------------------------------// 操作符/类注册// // 此文件不依赖 pybind11会被编译成 libidentity_conv_ops.so纯 C 二进制文件可以通过 dlopen 加载它而无需依赖 libpython。// ---------------------------------------------------------------------------TORCH_LIBRARY(my_ops, m){ // 注册 IdentityConvClass以便 Python 可以将其实例化为 torch.classes.my_ops.IdentityConvClass(channels)。 m.class_ (IdentityConvClass) .def(torch::init ()) .def(forward, IdentityConvClass::forward) .def(get_channels, IdentityConvClass::get_channels) // __obj_flatten__ 由 torch.export 的非严格跟踪器在切换到 FakeTensor 模式之前在 *真实* 的 C 对象上调用。 // 必须返回一个 (str, value) 对的元组以便 _check_valid_flat_script_obj 通过它会检查扁平序列中的每个元素是否为元组。我们将 channels_ 编码为单个命名条目没有张量叶子。 .def(__obj_flatten__, [](const c10::intrusive_ptr self) { return std::make_tuple( std::make_tuple(std::string(channels), self-channels_)); }) // def_pickle 启用 TorchScript 序列化。 .def_pickle( [](const c10::intrusive_ptr self) - int64_t { return self-channels_; }, [](int64_t channels) - c10::intrusive_ptr { return c10::make_intrusive (channels); }); // 注册 identity_conv_op 的模式。 m.def(identity_conv_op(Tensor x) - Tensor);}在 PyTorch 中使用自定义操作和类PyTorch 自定义类、函数及其在 C 中的注册会被构建成一个共享库libidentity_conv_ops.so可以使用 torch.ops.load_library 在 PyTorch 中加载和注册。为了兼容 torch.compile 和 torch.export还需要使用 register_fake_class 和 torch.library.register_fake 在 PyTorch 中注册自定义类和函数的“伪”抽象版本这样基于 FakeTensor 的符号跟踪就可以正确工作而无需在跟踪期间执行实际的 C/CUDA 代码。custom_ops.pycustom_ops.py加载 C / CUDA 共享库并设置 IdentityModel 使用的所有自定义 PyTorch 操作 1. torch.classes.my_ops.IdentityConvClass 由共享库注册 - 这里注册了一个伪/抽象版本以便 torch.export 可以跟踪持有该类实例的模块属性。 2. my_ops::identity_conv_op 模式 CPU CUDA 由共享库注册 - register_fake用于 torch.export / FakeTensor 的抽象实现。import osimport torchimport torch.library# ---------------------------------------------------------------------------# 1. 加载 C / CUDA 共享库。# 这会触发 TORCH_LIBRARY(my_ops, ...) 静态初始化器将 torch.classes.my_ops.IdentityConvClass 注册到 PyTorch 的全局操作符注册表中。# # 库路径可以通过 IDENTITY_CONV_OPS_LIB 环境变量覆盖默认情况下相对于此文件为 ../ext/libidentity_conv_ops.so。# ---------------------------------------------------------------------------_default_lib os.path.join( os.path.dirname(os.path.abspath(__file__)), .., ext, libidentity_conv_ops.so)_lib_path os.path.abspath( os.environ.get(IDENTITY_CONV_OPS_LIB, _default_lib))torch.ops.load_library(_lib_path)# ---------------------------------------------------------------------------# 2. 为 torch.export 跟踪注册 IdentityConvClass 的“伪”抽象版本。# # torch.export 使用基于 FakeTensor 的符号跟踪。当它遇到模块上的自定义类属性时会查找# • __obj_flatten__ - 为 pytree 扁平化返回 (leaves, context)# • __obj_unflatten__ - 从 (leaves, context) 重建对象# 这些由 register_fake_class 装饰的 Python 类提供。# ---------------------------------------------------------------------------from torch._library.fake_class_registry import register_fake_classregister_fake_class(my_ops::IdentityConvClass)class FakeIdentityConvClass: torch.export 期间使用的 IdentityConvClass 的抽象对应类。 def __init__(self, channels: int) - None: self.channels_ channels # -- torch.export 所需的 pytree 协议 ---------------------------- def __obj_flatten__(self): # 必须返回一个 (str, value) 对的元组与 C 的 __obj_flatten__ 返回的 ((channels, N),) 匹配。 return ((channels, self.channels_), ) classmethod def __obj_unflatten__(cls, flat): # flat 是 maybe_to_fake_obj 生成的可能是张量伪化的(key, value) 对序列。从中重建。 return cls(dict(flat)[channels]) # -- 抽象方法实现对 FakeTensors 进行操作 ------------ def forward(self, x: torch.Tensor) - torch.Tensor: # 形状/数据类型与输入一致 - 正确的抽象行为。 return torch.empty_like(x) def get_channels(self) - int: return self.channels_# ---------------------------------------------------------------------------# 3. 为 torch.export 跟踪注册 identity_conv_op 的伪抽象实现。# # 模式和两个实现CUDA 和 CPU已经通过 C 扩展使用 TORCH_LIBRARY / TORCH_LIBRARY_IMPL 进行了注册。Python 只需要提供抽象的形状/数据类型计算以便 torch.export 的 FakeTensor 解释器可以跟踪该操作。# ---------------------------------------------------------------------------torch.library.register_fake(my_ops::identity_conv_op)def _identity_conv_op_fake(x: torch.Tensor) - torch.Tensor: torch.export / FakeTensor 跟踪使用的抽象实现。 return torch.empty_like(x)# 方便的别名以便其他模块可以这样做from custom_ops import identity_conv_opidentity_conv_op torch.ops.my_ops.identity_conv_op加载共享库后可以使用 torch.classes 加载 PyTorch 自定义类使用 torch.ops 加载 PyTorch 自定义函数。model.pymodel.py定义 AOTInductor 演示中使用的四层 IdentityModel。层布局------------ layer1 : IdentityConv - 原生 PyTorch 操作符 layer2 : IdentityConvCustomClass - torch.classes C/CUDA 自定义类 layer3 : IdentityConvCustomOp - torch.library.custom_op C/CUDA 操作符 layer4 : IdentityConv - 原生 PyTorch 操作符每个层都是恒等变换因此对于任何输入 xmodel(x) x。import torchimport torch.nn as nn# 导入 custom_ops 会注册 C 扩展、伪类和自定义操作 - 必须在实例化任何模型之前进行。from custom_ops import identity_conv_op # noqa: F401# ---------------------------------------------------------------------------# 第 1/4 层 - 原生 PyTorch 逐通道 1×1 卷积恒等权重# ---------------------------------------------------------------------------class IdentityConv(nn.Module): 使用原生 PyTorch 操作符实现的恒等卷积。 使用内核大小为 1 且权重为 1.0 的逐通道 Conv2d相当于无操作输出 输入。该层开箱即用地兼容 torch.export 和 AOTInductor。 def __init__(self, channels: int) - None: super().__init__() self.conv nn.Conv2d( in_channelschannels, out_channelschannels, kernel_size(1, 1), stride(1, 1), padding(0, 0), dilation(1, 1), groupschannels, biasFalse, ) # 将所有权重设置为 1.0使卷积作为恒等操作。 self.conv.weight.data torch.ones(channels, 1, 1, 1) # 冻结权重 - 它们是常量不是可学习的参数。 self.conv.weight.requires_grad False def forward(self, x: torch.Tensor) - torch.Tensor: return self.conv(x)# ---------------------------------------------------------------------------# 第 2 层 - 通过 torch.classes 实现的自定义 C/CUDA 类# ---------------------------------------------------------------------------class IdentityConvCustomClass(nn.Module): 由 torch.classes C/CUDA 自定义类支持的恒等卷积。 在运行时前向调用会调度到 IdentityConvClasscsrc/identity_conv.cpp .cu内部注册的 CUDA 内核。 为了兼容 torch.export在 custom_ops.py 中通过 register_fake_class 注册了一个 FakeIdentityConvClass以便符号跟踪可以正常工作。 def __init__(self, channels: int) - None: super().__init__() self.obj torch.classes.my_ops.IdentityConvClass(channels) def forward(self, x: torch.Tensor) - torch.Tensor: return self.obj.forward(x)# ---------------------------------------------------------------------------# 第 3 层 - 通过 torch.library.custom_op 实现的自定义 C/CUDA 操作符# ---------------------------------------------------------------------------class IdentityConvCustomOp(nn.Module): 由 torch.library.custom_op C/CUDA 操作符支持的恒等卷积。 操作符my_ops::identity_conv_op在 custom_ops.py 中定义 • 一个用于 torch.export 跟踪的 register_fake 实现 • 一个调用 CUDA 内核的 register_kernel(cuda) 实现 def __init__(self, channels: int) - None: super().__init__() self.channels channels def forward(self, x: torch.Tensor) - torch.Tensor: return identity_conv_op(x)# ---------------------------------------------------------------------------# 完整模型# ---------------------------------------------------------------------------class IdentityModel(nn.Module): 用于 AOTInductor 演示的四层恒等模型。 def __init__(self, channels: int) - None: super().__init__() self.layer1 IdentityConv(channels) self.layer2 IdentityConvCustomClass(channels) self.layer3 IdentityConvCustomOp(channels) self.layer4 IdentityConv(channels) def forward(self, x: torch.Tensor) - torch.Tensor: x self.layer1(x) x self.layer2(x) x self.layer3(x) x self.layer4(x) return xdef create_model(channels: int 3) - IdentityModel: 返回一个在默认 CUDA 设备上处于评估模式的 IdentityModel。 return IdentityModel(channelschannels).cuda().eval()PyTorch 模型导出和降低如果为 torch.export 符号跟踪注册了所有自定义类和函数的伪抽象版本那么使用自定义类和自定义函数的 PyTorch 模型就可以通过 torch.export 导出。export_compile.py#!/usr/bin/env python3export_compile.py使用 torch.export 导出 IdentityModel并使用 torch._inductor.aoti_compile_and_package 进行编译。生成的包model.pt2会写入 artifacts/ 目录可以由 run_inference.pyPython和 C 推理二进制文件加载。用法从 python/ 目录运行 python export_compile.pyimport osimport sys# 确保 python/ 目录在路径中以便找到本地模块。sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))import torch# 导入 custom_ops 会加载 C 扩展并注册所有自定义操作。import custom_ops # noqa: F401from model import create_model# ---------------------------------------------------------------------------# 配置# ---------------------------------------------------------------------------CHANNELS 3BATCH_SIZE 1HEIGHT 224WIDTH 224# 将编译后的包保存在项目根目录的 artifacts/ 目录中。PACKAGE_PATH os.path.join(os.path.dirname(os.path.abspath(__file__)), .., artifacts, model.pt2)def main() - None: print( * 64) print(AOTInductor - 导出与编译) print( * 64) # ------------------------------------------------------------------ # 步骤 1实例化模型并验证即时执行的正确性 # ------------------------------------------------------------------ print(f\n[1/4] 创建 IdentityModel 通道数{CHANNELS}...) model create_model(channelsCHANNELS) x torch.randn(BATCH_SIZE, CHANNELS, HEIGHT, WIDTH, devicecuda, dtypetorch.float32) with torch.no_grad(): out model(x) assert torch.equal( x, out), (f导出前即时检查失败 f(最大差异 {(x - out).abs().max().item():.2e})) print( 即时验证通过 逐位相同) # ------------------------------------------------------------------ # 步骤 2使用 torch.export 导出 # ------------------------------------------------------------------ print(\n[2/4] 使用 torch.export.export() 导出模型 ...) with torch.no_grad(): exported_program torch.export.export(model, (x, )) print( 导出完成) print(f\n 导出的图:\n{exported_program.graph}) # ------------------------------------------------------------------ # 步骤 3使用 AOTInductor 编译 # ------------------------------------------------------------------ print( \n[3/4] 使用 torch._inductor.aoti_compile_and_package 进行编译 ...) package_path torch._inductor.aoti_compile_and_package( exported_program, package_pathPACKAGE_PATH, ) print(f 编译完成) print(f 包保存到: {os.path.abspath(package_path)}) # ------------------------------------------------------------------ # 步骤 4快速健全性检查 - 加载包并运行推理 # ------------------------------------------------------------------ print( \n[4/4] 快速健全性检查: 加载包并运行推理 ... ) compiled_model torch._inductor.aoti_load_package(package_path) with torch.no_grad(): out_compiled compiled_model(x) # aoti_load_package 返回一个可调用对象其输出是一个张量列表。 if isinstance(out_compiled, (list, tuple)): out_compiled out_compiled[0] assert torch.equal(x, out_compiled), ( f编译模型健全性检查失败 f(最大差异 {(x - out_compiled).abs().max().item():.2e})) print( 编译模型验证通过 逐位相同) print(\n * 64) print(f成功 包: {os.path.abspath(package_path)}) print( * 64)if __name__ __main__: main()从导出的图中可以看到自定义类 IdentityConvClass.forward 表示为对 torch.ops.higher_order.call_torchbind 的调用。自定义操作 identity_conv_op 表示为对 torch.ops.my_ops.identity_conv_op 的调用。graph(): %p_layer1_conv_weight : [num_users1] placeholder[targetp_layer1_conv_weight] %p_layer4_conv_weight : [num_users1] placeholder[targetp_layer4_conv_weight] %obj_layer2_obj : [num_users1] placeholder[targetobj_layer2_obj] %x : [num_users1] placeholder[targetx] %conv2d : [num_users1] call_function[targettorch.ops.aten.conv2d.default](args (%x, %p_layer1_conv_weight, None, [1, 1], [0, 0], [1, 1], 3), kwargs {}) %call_torchbind : [num_users1] call_function[targettorch.ops.higher_order.call_torchbind](args (%obj_layer2_obj, forward, %conv2d), kwargs {}) %identity_conv_op : [num_users1] call_function[targettorch.ops.my_ops.identity_conv_op.default](args (%call_torchbind,), kwargs {}) %conv2d_1 : [num_users1] call_function[targettorch.ops.aten.conv2d.default](args (%identity_conv_op, %p_layer4_conv_weight, None, [1, 1], [0, 0], [1, 1], 3), kwargs {}) return (conv2d_1,)导出的程序可以使用 torch._inductor.aoti_compile_and_package 进行编译和打包生成一个 model.pt2 包该包可以由 Python 和 C 推理程序加载。在执行编译后的模型时自定义类和自定义操作的实现将从共享库中加载并在运行时正确调度。run_inference.py#!/usr/bin/env python3run_inference.py加载 AOTInductor 编译的 IdentityModel 包model.pt2并运行推理以验证正确性。恒等模型的输出必须在严格的浮点容差范围内等于输入。用法在 export_compile.py 之后从 python/ 目录运行 python run_inference.py [MODEL_PATH [OP_LIB_PATH]]参数: MODEL_PATH 编译后的模型包.pt2的路径。 默认相对于此脚本为 ../artifacts/model.pt2。 OP_LIB_PATH 自定义操作共享库.so的路径。 提供时库路径会通过 IDENTITY_CONV_OPS_LIB 环境变量转发给 custom_ops.py以便 torch.ops.load_library 使用该文件而不是默认的 ../ext/libidentity_conv_ops.so。import osimport sys# 确保 python/ 目录在路径中以便找到本地模块。sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))import torchimport torch._inductor.codecache # 在 aoti_load_package 之前需要# ---------------------------------------------------------------------------# 解析命令行参数# ---------------------------------------------------------------------------_DEFAULT_PACKAGE_PATH os.path.join( os.path.dirname(os.path.abspath(__file__)), .., artifacts, model.pt2)PACKAGE_PATH sys.argv[1] if len(sys.argv) 1 else _DEFAULT_PACKAGE_PATHOP_LIB_PATH sys.argv[2] if len(sys.argv) 2 else None# ---------------------------------------------------------------------------# 如果提供了显式的库路径通过环境变量将其传递给 custom_ops.py以便 torch.ops.load_library 使用该文件。# ---------------------------------------------------------------------------if OP_LIB_PATH is not None: os.environ[IDENTITY_CONV_OPS_LIB] os.path.abspath(OP_LIB_PATH)# 导入 custom_ops 会在加载编译后的模型之前加载共享库并注册所有自定义操作。import custom_ops # noqa: F401# ---------------------------------------------------------------------------# 配置 - 必须与 export_compile.py 中使用的值匹配# ---------------------------------------------------------------------------CHANNELS 3BATCH_SIZE 1HEIGHT 224WIDTH 224def main() - None: print( * 64) print(AOTInductor - Python 推理) print( * 64) # ------------------------------------------------------------------ # 步骤 1加载编译后的模型包 # ------------------------------------------------------------------ pkg os.path.abspath(PACKAGE_PATH) if OP_LIB_PATH is not None: print(f 操作库 : {os.path.abspath(OP_LIB_PATH)}) print(f\n[1/3] 从以下路径加载编译后的模型:\n {pkg}) compiled_model torch._inductor.aoti_load_package(pkg) print( 模型加载成功。) # ------------------------------------------------------------------ # 步骤 2准备输入 # ------------------------------------------------------------------ x torch.randn(BATCH_SIZE, CHANNELS, HEIGHT, WIDTH, devicecuda, dtypetorch.float32) print(f\n[2/3] 输入 形状{list(x.shape)} 数据类型{x.dtype} f设备{x.device}) # ------------------------------------------------------------------ # 步骤 3运行推理并验证 # ------------------------------------------------------------------ print(\n[3/3] 运行推理 ...) with torch.no_grad(): output compiled_model(x) # aoti_load_package 返回一个可调用对象其输出是一个张量列表。 if isinstance(output, (list, tuple)): output output[0] print(f 输出形状{list(output.shape)} 数据类型{output.dtype}) if torch.equal(x, output): print(\n 验证通过 逐位相同) else: max_diff (x - output).abs().max().item() print(f\n 验证失败 最大差异 {max_diff} f — 期望输出逐位相同) sys.exit(1) print(\n * 64) print(成功 AOTInductor Python 推理验证通过。) print( * 64)if __name__ __main__: main()在纯 C 推理程序中可以使用 dlopen 加载和注册自定义类和自定义函数的共享库而无需任何 pybind11 或 libpython 依赖。run_inference.cpp/* * main.cpp *