Contrib ops

目录

contrib ops 域包含运行时默认内置的运算符。只有选定的运算符作为 contrib 运算符添加,以避免增加核心运行时包的二进制大小。在可能的情况下,应使用自定义运算符

Contrib 运算符列表

Contrib 运算符的 schema 在 ONNX Runtime 仓库中有所记录。

添加 Contrib 运算符

自定义运算符的 schema 和形状推断函数应使用 ONNX_CONTRIB_OPERATOR_SCHEMA 添加到 contrib_defs.cc 中。示例:Inverse 运算符

ONNX_CONTRIB_OPERATOR_SCHEMA(Inverse)
    .SetDomain(kMSDomain) // kMSDomain = "com.microsoft"
    .SinceVersion(1) // Same version used at op (symbolic) registration
    ...

新的运算符应具有完整的参考实现测试和形状推断测试。

参考实现 Python 测试应添加到 onnxruntime/test/python/contrib_ops 中。例如,aten_op_tests.py

形状推断 C++ 测试应添加到 onnxruntime/test/contrib_ops 中。例如,trilu_shape_inference_test.cc

运算符内核应使用 Compute 函数在 contrib 命名空间下实现,CPU 部分在 onnxruntime/contrib_ops/cpu/ 中,CUDA 部分在 onnxruntime/contrib_ops/cuda/ 中。

namespace onnxruntime {
namespace contrib {

class Inverse final : public OpKernel {
 public:
  explicit Inverse(const OpKernelInfo& info) : OpKernel(info) {}
  Status Compute(OpKernelContext* ctx) const override;

 private:
 ...
};

ONNX_OPERATOR_KERNEL_EX(
    Inverse,
    kMSDomain,
    1,
    kCpuExecutionProvider,
    KernelDefBuilder()
        .TypeConstraint("T", BuildKernelDefConstraints<float, double, MLFloat16>()),
    Inverse);

Status Inverse::Compute(OpKernelContext* ctx) const {
... // kernel implementation
}

}  // namespace contrib
}  // namespace onnxruntime

内核应在 cpu_contrib_kernels.cc 中注册用于 CPU,在 cuda_contrib_kernels.cc 中注册用于 CUDA。

现在您应该能够构建并安装 ONNX Runtime 以开始使用您的自定义运算符。

Contrib 运算符测试

测试应添加到 onnxruntime/test/contrib_ops/ 中。例如

namespace onnxruntime {
namespace test {

// Add a comprehensive set of unit tests for custom op kernel implementation

TEST(InverseContribOpTest, two_by_two_float) {
  OpTester test("Inverse", 1, kMSDomain); // custom opset version and domain
  test.AddInput<float>("X", {2, 2}, {4, 7, 2, 6});
  test.AddOutput<float>("Y", {2, 2}, {0.6f, -0.7f, -0.2f, 0.4f});
  test.Run();
}

...

}  // namespace test
}  // namespace onnxruntime