4、ONNX算子注册机制：OpSchema、OpSet、Domain的概念与关系
好，咱们今天来聊聊ONNX里最核心的一块——算子注册机制。说实话，我刚接触ONNX那会儿，也被OpSchema、OpSet、Domain这几个概念绕得有点晕。它们到底是个啥？互相之间怎么配合的？今天我就用我自己的理解，给你掰扯清楚。
4.1 为什么需要注册机制？
先问个问题：ONNX作为一个中间格式，它怎么知道一个算子长什么样？比如Conv，它有几个输入？几个输出？哪些属性是必须的？
答案就是——通过注册机制。每个算子在使用前，都得先在ONNX的“户口本”上登记。这个户口本，就是OpSchema。
核心思想：注册机制 = 给算子建立档案。档案里写清楚它的输入输出、属性、类型约束等所有信息。
我个人习惯把OpSchema理解成“算子的身份证”。你想想看，如果没有这张身份证，ONNX的runtime怎么知道这个算子该怎么处理？
4.2 OpSchema：算子的身份证
OpSchema是ONNX里描述一个算子所有元信息的类。它定义了：
- 算子名称：比如"Conv"、"Relu"
- 输入输出：数量、名称、类型约束
- 属性：哪些是必须的，哪些是可选的，默认值是什么
- 文档：算子的功能描述
- 类型约束：支持哪些数据类型（float、int等）
来看一个实际的例子。这是我之前在项目中写的一个自定义算子注册代码：
// 注册一个自定义算子 "MyCustomOp"
OpSchema my_schema("MyCustomOp", __FILE__, __LINE__);
// 定义输入：一个张量，名为 "input"
my_schema.Input(0, "input", "输入张量", "T");
// 定义输出：一个张量，名为 "output"
my_schema.Output(0, "output", "输出张量", "T");
// 定义属性：一个浮点数，名为 "scale"，默认值1.0
my_schema.Attr("scale", "缩放因子", AttributeProto::FLOAT, 1.0f);
// 定义类型约束：T 可以是 float 或 double
my_schema.TypeConstraint("T", {"tensor(float)", "tensor(double)"},
"float 或 double 类型");
// 设置算子所属的域和版本
my_schema.SinceVersion(1);
my_schema.SetDomain("my.custom.domain");嗯，这里要注意：SinceVersion 和 SetDomain 这两个调用特别关键。它们把OpSchema和OpSet、Domain联系起来了。
4.3 OpSet：算子的版本管理
OpSet，说白了就是一个“算子版本集合”。它把同一域下、同一版本的算子归到一起。
为什么会需要OpSet？因为算子会迭代啊。比如Conv在版本1时只支持2D卷积，版本11时支持了3D卷积。ONNX通过OpSet来管理这种演进。
我的经验：在ONNX官方标准中，每个OpSet都有一个版本号。比如ONNX opset 11就包含了当时所有官方算子的定义。你导出的模型如果用的是opset 11，那推理引擎也必须支持opset 11。
OpSet和OpSchema的关系是这样的：
- 每个OpSet包含多个OpSchema
- 每个OpSchema属于一个OpSet
- OpSet的版本号决定了算子定义的版本
举个例子：
// 创建一个OpSet，版本号为10，域为 "my.custom.domain"
OpSet my_opset("my.custom.domain", 10);
// 把之前定义的MyCustomOp注册到这个OpSet中
my_opset.Register(my_schema);
// 也可以直接链式注册
my_opset.Register(
OpSchema("AnotherOp", __FILE__, __LINE__)
.Input(0, "x", "输入", "T")
.Output(0, "y", "输出", "T")
.SinceVersion(1)
.SetDomain("my.custom.domain")
);4.4 Domain：命名空间与隔离
Domain，你可以把它理解成“命名空间”。它解决了算子命名冲突的问题。
ONNX官方算子都在 ai.onnx 这个域下。如果你自己写了一个自定义算子，也叫"Conv"，那就会和官方的冲突。怎么办？把你的算子放到你自己的域里，比如 my.company.conv。
| Domain | 说明 | 示例 |
|---|---|---|
ai.onnx |
ONNX官方标准算子域 | Conv, Relu, Gemm |
ai.onnx.ml |
机器学习相关算子域 | LabelEncoder, TreeEnsemble |
ai.onnx.training |
训练相关算子域 | Gradient, Momentum |
my.custom.domain |
自定义算子域 | MyCustomOp, MyFusion |
避坑指南：我曾经在项目中犯过一个错误——自定义算子忘了设置Domain，结果默认被归到了 ai.onnx 域下。推理时runtime报错说找不到这个算子。排查了半天才发现是Domain没设对。所以，自定义算子一定要显式设置Domain！
4.5 三者的关系：一张图说清楚
咱们来总结一下三者的关系：
- OpSchema：单个算子的定义（身份证）
- OpSet：同一域下、同一版本的算子集合（户口本）
- Domain：命名空间，隔离不同来源的算子（省份）
它们的关系可以这样理解：
Domain (省份)
└── OpSet v1 (户口本第一版)
├── OpSchema "Conv" (张三的身份证)
├── OpSchema "Relu" (李四的身份证)
└── OpSchema "MyCustomOp" (王五的身份证)
└── OpSet v2 (户口本第二版)
├── OpSchema "Conv" (张三更新后的身份证)
└── OpSchema "Gemm" (赵六的身份证)说白了，你要找一个算子，得先知道它在哪个Domain（省份），然后查对应版本的OpSet（户口本），最后找到具体的OpSchema（身份证）。
4.6 注册流程实战
最后，我给你展示一个完整的注册流程。这是我实际项目中用过的模式：
// 1. 定义OpSchema
OpSchema schema("MyAdd", __FILE__, __LINE__);
schema.Input(0, "a", "第一个输入", "T");
schema.Input(1, "b", "第二个输入", "T");
schema.Output(0, "c", "输出", "T");
schema.TypeConstraint("T", {"tensor(float)"}, "仅支持float");
schema.SinceVersion(1);
schema.SetDomain("my.custom.domain");
// 2. 创建OpSet并注册
OpSet opset("my.custom.domain", 1);
opset.Register(schema);
// 3. 将OpSet注册到全局注册表中
OpSchemaRegistry::RegisterOpSet(opset);
// 4. 验证注册是否成功
const auto* registered = OpSchemaRegistry::Schema("MyAdd", 1, "my.custom.domain");
if (registered) {
// 注册成功，可以正常使用了
std::cout << "算子注册成功！" << std::endl;
} else {
// 注册失败，检查原因
std::cerr << "算子注册失败！" << std::endl;
}我的建议：注册完成后，一定要做验证。别像我之前那样，注册完就以为万事大吉，结果推理时才发现根本没注册上。加个验证代码，花不了几秒钟，但能省下大把调试时间。
好了，关于OpSchema、OpSet、Domain的概念与关系，我就讲到这里。记住一句话：Domain是省份，OpSet是户口本，OpSchema是身份证。搞清楚了这三者的关系，自定义算子注册你就掌握了八成。下一节，咱们会深入讲讲如何编写一个完整的自定义算子，包括前向计算和反向梯度的实现。