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