ONNX Runtime 外部自定义算子 Schema 机制:--enable_external_custom_op_schemas 构建选项与完整可运行示例【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime本篇围绕 ONNX Runtime 仓库中 orttraining/orttraining/test/external_custom_ops 下的官方示例,讲解外部自定义算子 Schema(external custom op schemas)这一特性:它如何让用户把自定义算子的 Schema 注册(而非算子 kernel 本身)放在一个独立编译的 Python 扩展库中,运行时动态注入 ONNX 算子注册表。读完本文,你将掌握带--enable_external_custom_op_schemas构建 wheel 的方法、外部算子库的 Schema 定义与注册写法,以及从模型生成到推理验证的完整流程。一、特性定位:外部注册 ONNX 算子 Schema 的使用场景ONNX Runtime 通常要求模型中出现的算子(含自定义 domain 算子)要么属于已内置算子,要么通过 C 的RegisterCustomOps等 API 在进程内注册实现。而外部自定义算子 Schema针对的是另一类需求:算子定义(OpSchema)本身不在 onnxruntime 源码树里,而是由第三方以共享库形式提供,在运行时把 domain、算子 Schema、类型与形状推断、乃至函数体(function body)动态注册进 ONNX 注册表。仓库中的官方示例即演示这一 use-case,其 README(orttraining/orttraining/test/external_custom_ops/src/README.md)给出的原始信息如下:该示例要求 onnxruntime 以--enable_external_custom_op_schemas编译;该流程仅在 Ubuntu 上经过测试,Windows 上不可用;运行步骤(原文四步,完整保留):使用--build_wheel与--enable_external_custom_op_schemas构建 onnxruntime;在 Python 环境中安装该 onnxruntime wheel;在本文件夹下执行python3 -m pip install .安装自定义算子;执行python3 test.py进行测试。对应到 CMake,该特性由 cmake/CMakeLists.txt 中的选项控制,默认关闭:# Enable registering custom operator schemas from shared libraries in python environment. option(onnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS Enable registering user defined custom op schemas dynamically OFF)从选项注释看,其定位正是允许从 Python 环境中的共享库动态注册用户自定义的算子 Schema。二、构建前置:如何启用 --enable_external_custom_op_schemasREADME 第一步要求build.sh --build_wheel --enable_external_custom_op_schemas。这个参数在构建脚本中的定义与流转关系如下:参数定义在 tools/ci_build/build_args.py,且属于 Linux 专属参数组:parser.add_argument( --enable_external_custom_op_schemas, actionstore_true, helpEnable loading custom op schemas from external shared libraries (Ubuntu only)., )tools/ci_build/build.py 将其转换为 CMake 变量-Donnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMASON/OFF。CMake 侧在 cmake/CMakeLists.txt 中有三条硬性前置条件,不满足会直接FATAL_ERROR:if (onnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS) if (NOT CMAKE_SYSTEM_NAME STREQUAL Linux) message(FATAL_ERROR External custom operator schemas feature is only supported on Linux) endif() if (NOT ${ONNX_CUSTOM_PROTOC_EXECUTABLE} STREQUAL ) message(FATAL_ERROR External custom operator schemas is not supported with the user specified protoc executable) endif() if (NOT onnxruntime_ENABLE_TRAINING) message(FATAL_ERROR External custom operator schemas is supported only with --enable-training option) endif() add_custom_target(install_protobuf ALL DEPENDS ${PROTOBUF_LIB} protobuf::protoc) ... endif()即该特性:仅支持 Linux、不能与用户自定义的ONNX_CUSTOM_PROTOC_EXECUTABLE同时使用、且必须同时启用 onnxruntime 训练支持(onnxruntime_ENABLE_TRAINING)。因此完整的构建命令需要在--build_wheel --enable_external_custom_op_schemas之外再带上训练功能开关。同时,构建流程会额外安装 protobuf(cmake/CMakeLists.txt 中的install_protobuf目标执行cmake_install.cmake),这些产物最终会被打包进 wheel,供外部算子库编译时引用(见下节)。三、底层机制:这个构建选项到底改变了什么启用该选项后,构建流水线对 Python 绑定模块onnxruntime_pybind11_state做了四处关键处理,理解它们才能理解示例为什么能跑通。3.1 导出 onnx::* 与 protobuf::* 符号:version script 切换默认构建使用 onnxruntime/python/version_script.lds,只导出PyInit_onnxruntime_pybind11_state,其余符号全部 local(隐藏):VERS_1.0 { global: PyInit_onnxruntime_pybind11_state; # Hide everything else. local: *; };启用该特性后,cmake/onnxruntime_python.cmake 改用它旁边的version_script_expose_onnx_protobuf.lds:elseif(UNIX) if (onnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS) target_link_options(onnxruntime_pybind11_state PRIVATE LINKER:--version-script${ONNXRUNTIME_ROOT}/python/version_script_expose_onnx_protobuf.lds LINKER:--gc-sections) ...该脚本(onnxruntime/python/version_script_expose_onnx_protobuf.lds)额外把 ONNX 与 protobuf 的 C 符号全局导出:VERS_1.0 { global: PyInit_onnxruntime_pybind11_state; extern C { onnx::*; google::protobuf::*; }; # Hide everything else. local: *; };这样外部算子.so中对onnx::命名空间(如OpSchemaRegistry、RegisterSchema等)的引用,可以直接由 onnxruntime 已加载的进程内实例解析,而不是要求外部库再链接一份 libonnx。3.2 强制保留 onnx/protobuf 依赖:--no-as-needed / --as-needed仅靠 version script 还不够:链接器在--as-needed语义下可能裁掉只被符号导出引用的依赖。cmake/onnxruntime_python.cmake 通过精细插入链接选项解决:if (onnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS) ... if (ONNX_INDEX GREATER_EQUAL 0 AND PROTOBUF_INDEX GREATER_EQUAL 0) # Expect protobuf to follow onnx due to dependence list(INSERT onnxruntime_CUSTOM_EXTERNAL_LIBRARIES ${ONNX_INDEX} LINKER:--no-as-needed) list(INSERT onnxruntime_CUSTOM_EXTERNAL_LIBRARIES ${PROTOBUF_INDEX_NEXT} LINKER:--as-needed) ...即在 onnx 库前插入--no-as-needed(强制将其符号拉入)、在 protobuf 之后恢复--as-needed,确保 onnx 与 protobuf 的符号真实保留在.so内可供外部解析。3.3 dlopen 标志管理:RTLD_GLOBAL 与 RTLD_DEEPBINDcmake/onnxruntime_python.cmake 在启用该特性时,向模板 onnxruntime/python/_pybind_state.py.in 注入 dlopen 标志:if (onnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS) set(ONNXRUNTIME_SETDLOPENFLAGS_GLOBAL sys.setdlopenflags(os.RTLD_GLOBAL|os.RTLD_NOW|os.RTLD_DEEPBIND)) set(ONNXRUNTIME_SETDLOPENFLAGS_LOCAL sys.setdlopenflags(os.RTLD_LOCAL|os.RTLD_NOW|os.RTLD_DEEPBIND)) else() set(ONNXRUNTIME_SETDLOPENFLAGS_GLOBAL ) set(ONNXRUNTIME_SETDLOPENFLAGS_LOCAL ) endif()模板中对应的占位符包裹着核心绑定模块的导入:ONNXRUNTIME_IMPORT_PYTORCH_TO_RESOLVE_DLLS ONNXRUNTIME_SETDLOPENFLAGS_GLOBAL from .onnxruntime_pybind11_state import * # noqa ONNXRUNTIME_SETDLOPENFLAGS_LOCAL也就是说:导入onnxruntime_pybind11_state时先切到RTLD_GLOBAL|RTLD_NOW|RTLD_DEEPBIND(使其导出的 onnx/protobuf 符号对整个进程可见),导入完成后再切回RTLD_LOCAL|RTLD_NOW|RTLD_DEEPBIND。RTLD_DEEPBIND的作用,正如 test.py 的注释所解释的:Expose available (onnx::* and protobuf::*) symbols from onnxruntime to resolve references in the custom ops shared library. Deepbind flag is required to avoid conflicts with other instances of onnx/protobuf libraries.即:外部算子库引用的 onnx/protobuf 符号统一由 onnxruntime 导出,而 deepbind 避免与进程中其他 onnx/protobuf 实例(例如 pip 安装的 onnx 包)发生符号冲突。3.4 将 ONNX/protobuf 头文件随 wheel 分发外部算子库编译时需要 ONNX 头文件。启用该选项后,onnxruntime_pybind11_state的 POST_BUILD 阶段(cmake/onnxruntime_python.cmake)会在输出目录创建onnxruntime/external/include/,并分别建立指向构建产物中include/google(protobuf 头文件)与external/onnx/onnx(ONNX 头文件)的符号链接。随后,根目录 setup.py 在打包 wheel 时检测到onnxruntime/external目录存在(该目录只有启用此选项才会创建),会把其中全部文件收进 wheel:# Include files in onnxruntime/external if --enable_external_custom_op_schemas build.sh command # line option is specified. # If the options is not specified this following condition fails as onnxruntime/external folder is not created in the # build flow under the build binary directory. if path.isdir(path.join(onnxruntime, external)): extra.extend(...)这使得用户无需在本机单独克隆 ONNX 源码,即可用安装后的 wheel 内头文件编译外部算子扩展。四、示例算子 Foo:Schema 定义与运行时注册示例的核心文件是 orttraining/orttraining/test/external_custom_ops/src/foo_op.cpp,它定义了一个名为Foo的算子,行为是把输入张量原样复制到输出。4.1 算子 Schema 定义static const char FooDoc[] Foo copies input tensor to the output tensor.; ONNX_OPERATOR_SET_SCHEMA_EX( Foo, comExamples, com.examples, 1, false, OpSchema() .SetDoc(FooDoc) .Input(0, X, Input tensor, T) .Output(0, Y, Output tensor, T) .TypeConstraint( T, {tensor(float), tensor(int32), tensor(float16)}, Constrain input and output types to signed numeric tensors.) .TypeAndShapeInferenceFunction(FooShapeInference) .SetContextDependentFunctionBodyBuilder(...));要点(foo_op.cpp):使用ONNX_OPERATOR_SET_SCHEMA_EX宏(带显式 domain 字符串参数的变体)注册:domain 为com.examples,算子版本 1;输入X与输出Y共享类型约束T,限定为tensor(float)、tensor(int32)、tensor(float16)三种数值张量;形状推断函数FooShapeInference调用 ONNX 辅助函数propagateShapeAndTypeFromFirstInput(ctx),即输出直接继承第一个输入的 shape 与类型;通过SetContextDependentFunctionBodyBuilder注册上下文相关的函数体构造器:当运行时能提供输入类型时(ctx.getInputType(0) ! nullptr),用FunctionBodyHelper::BuildNodes生成一个单节点函数体{Y, Identity, {X}},即把Foo展开为Identity;若拿不到输入类型则返回false,表示无法构造正确的函数体。4.2 库加载即注册:静态初始化器 空 pybind11 模块static bool registerOps() { auto d OpSchemaRegistry::DomainToVersionRange::Instance(); d.AddDomainToVersion(com.examples, 1, 1); auto schema GetOpSchemaONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(comExamples, 1, Foo)(); RegisterSchema(schema); std::cerr Successfully registered custom op std::endl; return true; } static bool ret registerOps(); PYBIND11_MODULE(orttraining_external_custom_ops, m) { }(foo_op.cpp)这里有两个关键手法:static bool ret registerOps();借助全局静态对象的初始化时机,在.so被dlopen时自动执行注册:先把 domaincom.examples的版本区间(1..1)加入OpSchemaRegistry::DomainToVersionRange,再注册具体算子 Schema,并在 stderr 打印Successfully registered custom op(这也是测试运行时的预期输出);PYBIND11_MODULE(orttraining_external_custom_ops, m) {}是一个空模块——它不提供任何 Python API,唯一作用是生成PyInit_orttraining_external_custom_ops入口,让这个.so能被 Python 正常import,从而触发上面的静态注册。这正是外部 Schema 注入模式:用户程序只要import该扩展模块,com.examples.Foo就进入进程级 ONNX 算子注册表,后续加载包含该算子的模型即可通过 Schema 校验与类型/形状推断。五、编译自定义算子库:setup.py 与 CMakeLists.txt示例目录提供了一套以 CMake 为后端的 pip 安装流程,即 README 第三步的python3 -m pip install .。setup.py 自定义了build_ext,在构建扩展时向 CMake 传入四个关键变量:subprocess.check_call( [ cmake, f-DPYBIND11_PYTHON_VERSION{...}, f-Dpybind11_DIR{pybind11.get_cmake_dir()}, f-DONNX_INCLUDE{os.path.dirname(os.path.dirname(onnx.__file__))}, -DONNXRUNTIME_EXTERNAL_INCLUDE{}.format( os.path.join(os.path.join(os.path.dirname(onnxruntime.__file__), external), include)), f-DCMAKE_LIBRARY_OUTPUT_DIRECTORY{extdir}, ext.sourcedir, ], cwdself.build_temp, )ONNX_INCLUDE指向 pip 安装的 onnx 包目录,提供 ONNX 头文件;ONNXRUNTIME_EXTERNAL_INCLUDE指向已安装的 onnxruntime wheel 内onnxruntime/external/include——即第三节 3.4 中打包进 wheel 的那份头文件。这印证了该构建选项对 wheel 打包的改造正是为了支持离线编译外部算子库。CMakeLists.txt 的关键配置:add_compile_definitions(ONNX_ML1) add_compile_definitions(ONNX_NAMESPACEonnx) add_compile_definitions(ONNX_USER_LITE_PROTOON) ... pybind11_add_module(orttraining_external_custom_ops src/foo_op.cpp)从这三条编译定义看,外部算子库被刻意约束到与 onnxruntime 内部一致的 ONNX 构建模式:ONNX_NAMESPACEonnx(统一命名空间,与 version script 导出的onnx::*符号对应)、ONNX_USER_LITE_PROTOON(lite protobuf 模式)、ONNX_ML1。从源码结构看,这保证了外部库在链接期生成的符号引用与运行时由 onnxruntime 解析的符号完全一致,不会出现命名空间或 proto 布局错配。此外,若未显式传入ONNXRUNTIME_EXTERNAL_INCLUDE/ONNX_INCLUDE,CMake 会通过python -c import onnxruntime ...等方式自动探测(CMakeLists.txt)。六、模型构建与测试验证6.1 生成测试模型testdata/gen_model.py 用 ONNX helper 构造了一个只包含一个Foo节点的最小模型:graph helper.make_graph( [helper.make_node(Foo, [input1], [output1], , , com.examples)], external_custom_op_example_model, [helper.make_tensor_value_info(input1, helper.TensorProto.FLOAT, [2, 2])], [helper.make_tensor_value_info(output1, helper.TensorProto.FLOAT, [2, 2])], [], ) model helper.make_model(graph) opset model.opset_import.add() opset.version 1 opset.domain com.examples注意模型显式声明了自定义 domaincom.examples的 opset(version 1)——这与foo_op.cpp注册的 domain/版本严格对应,否则图校验会因找不到该 domain 的 opset 而失败。仓库中已附带生成好的 testdata/model.onnx,测试可直接使用。6.2 运行测试test.py 完整逻辑如下:import numpy as np # Restore dlopen flags. import orttraining_external_custom_ops # noqa: F401 # Expose available (onnx::* and protobuf::*) symbols from onnxruntime to resolve references in # the custom ops shared library. Deepbind flag is required to avoid conflicts with other # instances of onnx/protobuf libraries. import onnxruntime so onnxruntime.SessionOptions() sess onnxruntime.InferenceSession(testdata/model.onnx, so) input np.random.rand(2, 2).astype(np.float32) output sess.run(None, {input1: input})[0] np.testing.assert_equal(input, output)流程解读:import orttraining_external_custom_ops加载外部算子扩展,静态初始化器随即把com.examples.Foo的 Schema 注册进进程,stderr 输出Successfully registered custom op;导入 onnxruntime 后,按 3.1/3.3 节的机制,其导出的onnx::*、google::protobuf::*符号可供外部算子库解析引用;用InferenceSession加载testdata/model.onnx并跑一次 2x2 float32 推理;由于Foo的函数体等价于Identity,测试断言input与output逐元素相等。至此,README 的第四步python3 test.py即完成端到端验证:外部 Schema 注册 → 模型校验通过 → 推理结果正确。七、适用范围与注意事项结合 README、构建参数与 CMake 约束,使用此特性需要同时满足:约束依据仅在 Linux(官方在 Ubuntu 上测试),Windows 不可用README;build_args.py 参数 help Ubuntu only;CMakeLists.txtFATAL_ERROR必须自编译 wheel(--build_wheel--enable_external_custom_op_schemas),普通安装渠道的 wheel 不具备导出的 onnx/protobuf 符号与external/include头文件README 第 1、2 步;setup.py 中onnxruntime/external目录仅在启用该选项时创建必须同时启用 onnxruntime 训练支持(onnxruntime_ENABLE_TRAINING)CMakeLists.txt不能与用户自定义的 protoc 可执行文件(ONNX_CUSTOM_PROTOC_EXECUTABLE)组合CMakeLists.txt外部算子库需与 ORT 使用一致的 ONNX 编译模式(命名空间onnx、lite proto),头文件可取自 wheel 内onnxruntime/external/include或 pip onnx 包示例 CMakeLists.txt;示例 setup.py模型侧需声明自定义 domain 的 opset,且算子的 domain/版本与注册代码一致gen_model.py该示例代码量很小(一个约 60 行的 C 文件加一套 CMake/pip 脚手架),却是理解 ONNX Runtime 如何在运行时接纳用户自定义算子定义的完整参考:version script 决定符号边界,链接选项决定依赖保留,_pybind_state.py的 dlopen 标志决定符号可见性与隔离方式,而ONNX_OPERATOR_SET_SCHEMA_EX加静态注册则构成外部 Schema 注入的最小实现范式。若要扩展自己的外部算子库,照此目录结构替换 domain、算子名、类型约束与函数体构造器,并保证编译定义与 wheel 内 ONNX 头文件一致,即可复用同一套构建与验证流程。【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/13 3:25:42 · 阅读 1.2k ONNX Runtime 外部自定义算子 Schema 机制:--enable_external_custom_op_schemas 构建选项与完整可运行示例
ONNX Runtime 外部自定义算子 Schema 机制:--enable_external_custom_op_schemas 构建选项与完整可运行示例【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime本篇围绕 ONNX Runtime 仓库中 orttraining/orttraining/test/external_custom_ops 下的官方示例,讲解外部自定义算子 Schema(external custom op schemas)这一特性:它如何让用户把自定义算子的 Schema 注册(而非算子 kernel 本身)放在一个独立编译的 Python 扩展库中,运行时动态注入 ONNX 算子注册表。读完本文,你将掌握带--enable_external_custom_op_schemas构建 wheel 的方法、外部算子库的 Schema 定义与注册写法,以及从模型生成到推理验证的完整流程。一、特性定位:外部注册 ONNX 算子 Schema 的使用场景ONNX Runtime 通常要求模型中出现的算子(含自定义 domain 算子)要么属于已内置算子,要么通过 C 的RegisterCustomOps等 API 在进程内注册实现。而外部自定义算子 Schema针对的是另一类需求:算子定义(OpSchema)本身不在 onnxruntime 源码树里,而是由第三方以共享库形式提供,在运行时把 domain、算子 Schema、类型与形状推断、乃至函数体(function body)动态注册进 ONNX 注册表。仓库中的官方示例即演示这一 use-case,其 README(orttraining/orttraining/test/external_custom_ops/src/README.md)给出的原始信息如下:该示例要求 onnxruntime 以--enable_external_custom_op_schemas编译;该流程仅在 Ubuntu 上经过测试,Windows 上不可用;运行步骤(原文四步,完整保留):使用--build_wheel与--enable_external_custom_op_schemas构建 onnxruntime;在 Python 环境中安装该 onnxruntime wheel;在本文件夹下执行python3 -m pip install .安装自定义算子;执行python3 test.py进行测试。对应到 CMake,该特性由 cmake/CMakeLists.txt 中的选项控制,默认关闭:# Enable registering custom operator schemas from shared libraries in python environment. option(onnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS Enable registering user defined custom op schemas dynamically OFF)从选项注释看,其定位正是允许从 Python 环境中的共享库动态注册用户自定义的算子 Schema。二、构建前置:如何启用 --enable_external_custom_op_schemasREADME 第一步要求build.sh --build_wheel --enable_external_custom_op_schemas。这个参数在构建脚本中的定义与流转关系如下:参数定义在 tools/ci_build/build_args.py,且属于 Linux 专属参数组:parser.add_argument( --enable_external_custom_op_schemas, actionstore_true, helpEnable loading custom op schemas from external shared libraries (Ubuntu only)., )tools/ci_build/build.py 将其转换为 CMake 变量-Donnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMASON/OFF。CMake 侧在 cmake/CMakeLists.txt 中有三条硬性前置条件,不满足会直接FATAL_ERROR:if (onnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS) if (NOT CMAKE_SYSTEM_NAME STREQUAL Linux) message(FATAL_ERROR External custom operator schemas feature is only supported on Linux) endif() if (NOT ${ONNX_CUSTOM_PROTOC_EXECUTABLE} STREQUAL ) message(FATAL_ERROR External custom operator schemas is not supported with the user specified protoc executable) endif() if (NOT onnxruntime_ENABLE_TRAINING) message(FATAL_ERROR External custom operator schemas is supported only with --enable-training option) endif() add_custom_target(install_protobuf ALL DEPENDS ${PROTOBUF_LIB} protobuf::protoc) ... endif()即该特性:仅支持 Linux、不能与用户自定义的ONNX_CUSTOM_PROTOC_EXECUTABLE同时使用、且必须同时启用 onnxruntime 训练支持(onnxruntime_ENABLE_TRAINING)。因此完整的构建命令需要在--build_wheel --enable_external_custom_op_schemas之外再带上训练功能开关。同时,构建流程会额外安装 protobuf(cmake/CMakeLists.txt 中的install_protobuf目标执行cmake_install.cmake),这些产物最终会被打包进 wheel,供外部算子库编译时引用(见下节)。三、底层机制:这个构建选项到底改变了什么启用该选项后,构建流水线对 Python 绑定模块onnxruntime_pybind11_state做了四处关键处理,理解它们才能理解示例为什么能跑通。3.1 导出 onnx::* 与 protobuf::* 符号:version script 切换默认构建使用 onnxruntime/python/version_script.lds,只导出PyInit_onnxruntime_pybind11_state,其余符号全部 local(隐藏):VERS_1.0 { global: PyInit_onnxruntime_pybind11_state; # Hide everything else. local: *; };启用该特性后,cmake/onnxruntime_python.cmake 改用它旁边的version_script_expose_onnx_protobuf.lds:elseif(UNIX) if (onnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS) target_link_options(onnxruntime_pybind11_state PRIVATE LINKER:--version-script${ONNXRUNTIME_ROOT}/python/version_script_expose_onnx_protobuf.lds LINKER:--gc-sections) ...该脚本(onnxruntime/python/version_script_expose_onnx_protobuf.lds)额外把 ONNX 与 protobuf 的 C 符号全局导出:VERS_1.0 { global: PyInit_onnxruntime_pybind11_state; extern C { onnx::*; google::protobuf::*; }; # Hide everything else. local: *; };这样外部算子.so中对onnx::命名空间(如OpSchemaRegistry、RegisterSchema等)的引用,可以直接由 onnxruntime 已加载的进程内实例解析,而不是要求外部库再链接一份 libonnx。3.2 强制保留 onnx/protobuf 依赖:--no-as-needed / --as-needed仅靠 version script 还不够:链接器在--as-needed语义下可能裁掉只被符号导出引用的依赖。cmake/onnxruntime_python.cmake 通过精细插入链接选项解决:if (onnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS) ... if (ONNX_INDEX GREATER_EQUAL 0 AND PROTOBUF_INDEX GREATER_EQUAL 0) # Expect protobuf to follow onnx due to dependence list(INSERT onnxruntime_CUSTOM_EXTERNAL_LIBRARIES ${ONNX_INDEX} LINKER:--no-as-needed) list(INSERT onnxruntime_CUSTOM_EXTERNAL_LIBRARIES ${PROTOBUF_INDEX_NEXT} LINKER:--as-needed) ...即在 onnx 库前插入--no-as-needed(强制将其符号拉入)、在 protobuf 之后恢复--as-needed,确保 onnx 与 protobuf 的符号真实保留在.so内可供外部解析。3.3 dlopen 标志管理:RTLD_GLOBAL 与 RTLD_DEEPBINDcmake/onnxruntime_python.cmake 在启用该特性时,向模板 onnxruntime/python/_pybind_state.py.in 注入 dlopen 标志:if (onnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS) set(ONNXRUNTIME_SETDLOPENFLAGS_GLOBAL sys.setdlopenflags(os.RTLD_GLOBAL|os.RTLD_NOW|os.RTLD_DEEPBIND)) set(ONNXRUNTIME_SETDLOPENFLAGS_LOCAL sys.setdlopenflags(os.RTLD_LOCAL|os.RTLD_NOW|os.RTLD_DEEPBIND)) else() set(ONNXRUNTIME_SETDLOPENFLAGS_GLOBAL ) set(ONNXRUNTIME_SETDLOPENFLAGS_LOCAL ) endif()模板中对应的占位符包裹着核心绑定模块的导入:ONNXRUNTIME_IMPORT_PYTORCH_TO_RESOLVE_DLLS ONNXRUNTIME_SETDLOPENFLAGS_GLOBAL from .onnxruntime_pybind11_state import * # noqa ONNXRUNTIME_SETDLOPENFLAGS_LOCAL也就是说:导入onnxruntime_pybind11_state时先切到RTLD_GLOBAL|RTLD_NOW|RTLD_DEEPBIND(使其导出的 onnx/protobuf 符号对整个进程可见),导入完成后再切回RTLD_LOCAL|RTLD_NOW|RTLD_DEEPBIND。RTLD_DEEPBIND的作用,正如 test.py 的注释所解释的:Expose available (onnx::* and protobuf::*) symbols from onnxruntime to resolve references in the custom ops shared library. Deepbind flag is required to avoid conflicts with other instances of onnx/protobuf libraries.即:外部算子库引用的 onnx/protobuf 符号统一由 onnxruntime 导出,而 deepbind 避免与进程中其他 onnx/protobuf 实例(例如 pip 安装的 onnx 包)发生符号冲突。3.4 将 ONNX/protobuf 头文件随 wheel 分发外部算子库编译时需要 ONNX 头文件。启用该选项后,onnxruntime_pybind11_state的 POST_BUILD 阶段(cmake/onnxruntime_python.cmake)会在输出目录创建onnxruntime/external/include/,并分别建立指向构建产物中include/google(protobuf 头文件)与external/onnx/onnx(ONNX 头文件)的符号链接。随后,根目录 setup.py 在打包 wheel 时检测到onnxruntime/external目录存在(该目录只有启用此选项才会创建),会把其中全部文件收进 wheel:# Include files in onnxruntime/external if --enable_external_custom_op_schemas build.sh command # line option is specified. # If the options is not specified this following condition fails as onnxruntime/external folder is not created in the # build flow under the build binary directory. if path.isdir(path.join(onnxruntime, external)): extra.extend(...)这使得用户无需在本机单独克隆 ONNX 源码,即可用安装后的 wheel 内头文件编译外部算子扩展。四、示例算子 Foo:Schema 定义与运行时注册示例的核心文件是 orttraining/orttraining/test/external_custom_ops/src/foo_op.cpp,它定义了一个名为Foo的算子,行为是把输入张量原样复制到输出。4.1 算子 Schema 定义static const char FooDoc[] Foo copies input tensor to the output tensor.; ONNX_OPERATOR_SET_SCHEMA_EX( Foo, comExamples, com.examples, 1, false, OpSchema() .SetDoc(FooDoc) .Input(0, X, Input tensor, T) .Output(0, Y, Output tensor, T) .TypeConstraint( T, {tensor(float), tensor(int32), tensor(float16)}, Constrain input and output types to signed numeric tensors.) .TypeAndShapeInferenceFunction(FooShapeInference) .SetContextDependentFunctionBodyBuilder(...));要点(foo_op.cpp):使用ONNX_OPERATOR_SET_SCHEMA_EX宏(带显式 domain 字符串参数的变体)注册:domain 为com.examples,算子版本 1;输入X与输出Y共享类型约束T,限定为tensor(float)、tensor(int32)、tensor(float16)三种数值张量;形状推断函数FooShapeInference调用 ONNX 辅助函数propagateShapeAndTypeFromFirstInput(ctx),即输出直接继承第一个输入的 shape 与类型;通过SetContextDependentFunctionBodyBuilder注册上下文相关的函数体构造器:当运行时能提供输入类型时(ctx.getInputType(0) ! nullptr),用FunctionBodyHelper::BuildNodes生成一个单节点函数体{Y, Identity, {X}},即把Foo展开为Identity;若拿不到输入类型则返回false,表示无法构造正确的函数体。4.2 库加载即注册:静态初始化器 空 pybind11 模块static bool registerOps() { auto d OpSchemaRegistry::DomainToVersionRange::Instance(); d.AddDomainToVersion(com.examples, 1, 1); auto schema GetOpSchemaONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(comExamples, 1, Foo)(); RegisterSchema(schema); std::cerr Successfully registered custom op std::endl; return true; } static bool ret registerOps(); PYBIND11_MODULE(orttraining_external_custom_ops, m) { }(foo_op.cpp)这里有两个关键手法:static bool ret registerOps();借助全局静态对象的初始化时机,在.so被dlopen时自动执行注册:先把 domaincom.examples的版本区间(1..1)加入OpSchemaRegistry::DomainToVersionRange,再注册具体算子 Schema,并在 stderr 打印Successfully registered custom op(这也是测试运行时的预期输出);PYBIND11_MODULE(orttraining_external_custom_ops, m) {}是一个空模块——它不提供任何 Python API,唯一作用是生成PyInit_orttraining_external_custom_ops入口,让这个.so能被 Python 正常import,从而触发上面的静态注册。这正是外部 Schema 注入模式:用户程序只要import该扩展模块,com.examples.Foo就进入进程级 ONNX 算子注册表,后续加载包含该算子的模型即可通过 Schema 校验与类型/形状推断。五、编译自定义算子库:setup.py 与 CMakeLists.txt示例目录提供了一套以 CMake 为后端的 pip 安装流程,即 README 第三步的python3 -m pip install .。setup.py 自定义了build_ext,在构建扩展时向 CMake 传入四个关键变量:subprocess.check_call( [ cmake, f-DPYBIND11_PYTHON_VERSION{...}, f-Dpybind11_DIR{pybind11.get_cmake_dir()}, f-DONNX_INCLUDE{os.path.dirname(os.path.dirname(onnx.__file__))}, -DONNXRUNTIME_EXTERNAL_INCLUDE{}.format( os.path.join(os.path.join(os.path.dirname(onnxruntime.__file__), external), include)), f-DCMAKE_LIBRARY_OUTPUT_DIRECTORY{extdir}, ext.sourcedir, ], cwdself.build_temp, )ONNX_INCLUDE指向 pip 安装的 onnx 包目录,提供 ONNX 头文件;ONNXRUNTIME_EXTERNAL_INCLUDE指向已安装的 onnxruntime wheel 内onnxruntime/external/include——即第三节 3.4 中打包进 wheel 的那份头文件。这印证了该构建选项对 wheel 打包的改造正是为了支持离线编译外部算子库。CMakeLists.txt 的关键配置:add_compile_definitions(ONNX_ML1) add_compile_definitions(ONNX_NAMESPACEonnx) add_compile_definitions(ONNX_USER_LITE_PROTOON) ... pybind11_add_module(orttraining_external_custom_ops src/foo_op.cpp)从这三条编译定义看,外部算子库被刻意约束到与 onnxruntime 内部一致的 ONNX 构建模式:ONNX_NAMESPACEonnx(统一命名空间,与 version script 导出的onnx::*符号对应)、ONNX_USER_LITE_PROTOON(lite protobuf 模式)、ONNX_ML1。从源码结构看,这保证了外部库在链接期生成的符号引用与运行时由 onnxruntime 解析的符号完全一致,不会出现命名空间或 proto 布局错配。此外,若未显式传入ONNXRUNTIME_EXTERNAL_INCLUDE/ONNX_INCLUDE,CMake 会通过python -c import onnxruntime ...等方式自动探测(CMakeLists.txt)。六、模型构建与测试验证6.1 生成测试模型testdata/gen_model.py 用 ONNX helper 构造了一个只包含一个Foo节点的最小模型:graph helper.make_graph( [helper.make_node(Foo, [input1], [output1], , , com.examples)], external_custom_op_example_model, [helper.make_tensor_value_info(input1, helper.TensorProto.FLOAT, [2, 2])], [helper.make_tensor_value_info(output1, helper.TensorProto.FLOAT, [2, 2])], [], ) model helper.make_model(graph) opset model.opset_import.add() opset.version 1 opset.domain com.examples注意模型显式声明了自定义 domaincom.examples的 opset(version 1)——这与foo_op.cpp注册的 domain/版本严格对应,否则图校验会因找不到该 domain 的 opset 而失败。仓库中已附带生成好的 testdata/model.onnx,测试可直接使用。6.2 运行测试test.py 完整逻辑如下:import numpy as np # Restore dlopen flags. import orttraining_external_custom_ops # noqa: F401 # Expose available (onnx::* and protobuf::*) symbols from onnxruntime to resolve references in # the custom ops shared library. Deepbind flag is required to avoid conflicts with other # instances of onnx/protobuf libraries. import onnxruntime so onnxruntime.SessionOptions() sess onnxruntime.InferenceSession(testdata/model.onnx, so) input np.random.rand(2, 2).astype(np.float32) output sess.run(None, {input1: input})[0] np.testing.assert_equal(input, output)流程解读:import orttraining_external_custom_ops加载外部算子扩展,静态初始化器随即把com.examples.Foo的 Schema 注册进进程,stderr 输出Successfully registered custom op;导入 onnxruntime 后,按 3.1/3.3 节的机制,其导出的onnx::*、google::protobuf::*符号可供外部算子库解析引用;用InferenceSession加载testdata/model.onnx并跑一次 2x2 float32 推理;由于Foo的函数体等价于Identity,测试断言input与output逐元素相等。至此,README 的第四步python3 test.py即完成端到端验证:外部 Schema 注册 → 模型校验通过 → 推理结果正确。七、适用范围与注意事项结合 README、构建参数与 CMake 约束,使用此特性需要同时满足:约束依据仅在 Linux(官方在 Ubuntu 上测试),Windows 不可用README;build_args.py 参数 help Ubuntu only;CMakeLists.txtFATAL_ERROR必须自编译 wheel(--build_wheel--enable_external_custom_op_schemas),普通安装渠道的 wheel 不具备导出的 onnx/protobuf 符号与external/include头文件README 第 1、2 步;setup.py 中onnxruntime/external目录仅在启用该选项时创建必须同时启用 onnxruntime 训练支持(onnxruntime_ENABLE_TRAINING)CMakeLists.txt不能与用户自定义的 protoc 可执行文件(ONNX_CUSTOM_PROTOC_EXECUTABLE)组合CMakeLists.txt外部算子库需与 ORT 使用一致的 ONNX 编译模式(命名空间onnx、lite proto),头文件可取自 wheel 内onnxruntime/external/include或 pip onnx 包示例 CMakeLists.txt;示例 setup.py模型侧需声明自定义 domain 的 opset,且算子的 domain/版本与注册代码一致gen_model.py该示例代码量很小(一个约 60 行的 C 文件加一套 CMake/pip 脚手架),却是理解 ONNX Runtime 如何在运行时接纳用户自定义算子定义的完整参考:version script 决定符号边界,链接选项决定依赖保留,_pybind_state.py的 dlopen 标志决定符号可见性与隔离方式,而ONNX_OPERATOR_SET_SCHEMA_EX加静态注册则构成外部 Schema 注入的最小实现范式。若要扩展自己的外部算子库,照此目录结构替换 domain、算子名、类型约束与函数体构造器,并保证编译定义与 wheel 内 ONNX 头文件一致,即可复用同一套构建与验证流程。【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考