Benutzerdefinierte Operatoren

ONNX Runtime bietet Optionen zur Ausführung benutzerdefinierter Operatoren, die keine offiziellen ONNX-Operatoren sind. Beachten Sie, dass sich benutzerdefinierte Operatoren von Contrib-Ops unterscheiden, die ausgewählte inoffizielle ONNX-Operatoren sind, die direkt in ORT integriert sind.

Inhalt

Einen benutzerdefinierten Operator definieren und registrieren

Seit onnxruntime 1.16 kann ein benutzerdefinierter Operator einfach als Funktion implementiert werden

void KernelOne(const Ort::Custom::Tensor<float>& X,
               const Ort::Custom::Tensor<float>& Y,
               Ort::Custom::Tensor<float>& Z) {
  auto input_shape = X.Shape();
  auto x_raw = X.Data();
  auto y_raw = Y.Data();
  auto z_raw = Z.Allocate(input_shape);
  for (int64_t i = 0; i < Z.NumberOfElement(); ++i) {
    z_raw[i] = x_raw[i] + y_raw[i];
  }
}

int main() {
  Ort::CustomOpDomain v1_domain{"v1"};
  // please make sure that custom_op_one has the same lifetime as the consuming session
  std::unique_ptr<OrtLiteCustomOp> custom_op_one{Ort::Custom::CreateLiteCustomOp("CustomOpOne", "CPUExecutionProvider", KernelOne)};
  v1_domain.Add(custom_op_one.get());
  Ort::SessionOptions session_options;
  session_options.Add(v1_domain);
  // create a session with the session_options ...
}

Für benutzerdefinierte Operatoren mit Attributen werden auch Strukturen unterstützt:

struct Merge {
  Merge(const OrtApi* ort_api, const OrtKernelInfo* info) {
    int64_t reverse;
    ORT_ENFORCE(ort_api->KernelInfoGetAttribute_int64(info, "reverse", &reverse) == nullptr);
    reverse_ = reverse != 0;
  }
  // a "Compute" member function is required to be present
  void Compute(const Ort::Custom::Tensor<std::string_view>& strings_in,
               std::string_view string_in,
               Ort::Custom::Tensor<std::string>* strings_out) {
    std::vector<std::string> string_pool;
    for (const auto& s : strings_in.Data()) {
      string_pool.emplace_back(s.data(), s.size());
    }
    string_pool.emplace_back(string_in.data(), string_in.size());
    if (reverse_) {
      for (auto& str : string_pool) {
        std::reverse(str.begin(), str.end());
      }
      std::reverse(string_pool.begin(), string_pool.end());
    }
    strings_out->SetStringOutput(string_pool, {static_cast<int64_t>(string_pool.size())});
  }
  bool reverse_ = false;
};

int main() {
  Ort::CustomOpDomain v2_domain{"v2"};
  // please make sure that mrg_op_ptr has the same lifetime as the consuming session
  std::unique_ptr<Ort::Custom::OrtLiteCustomOp> mrg_op_ptr{Ort::Custom::CreateLiteCustomOp<Merge>("Merge", "CPUExecutionProvider")};
  v2_domain.Add(mrg_op_ptr.get());
  Ort::SessionOptions session_options;
  session_options.Add(v2_domain);
  // create a session with the session_options ...
}

Eine „Compute“-Memberfunktion ist für die Struktur erforderlich, um als benutzerdefinierter Operator ausgeführt zu werden.

Für beide Fälle

  • Eingaben müssen als const-Referenzen deklariert werden.
  • Ausgaben müssen als non-const-Referenzen deklariert werden.
  • Ort::Custom::Tensor::Shape() gibt die Eingabeform zurück.
  • Ort::Custom::Tensor::Data() gibt Rohdaten der Eingabe zurück.
  • Ort::Custom::Tensor::NumberOfElement() gibt die Anzahl der Elemente in der Eingabe zurück.
  • Ort::Custom::Tensor::Allocate(…) reserviert eine Ausgabe und gibt die Adresse der Rohdaten zurück.
  • Unterstützte Template-Argumente sind: int8_t, int16_t, int32_t, int64_t, float, double.
  • Unterstützt std::string_view als Eingabe und std::string als Ausgabe. Die Verwendung finden Sie hier.
  • Für benutzerdefinierte Operatorfunktionen, die auf CPUExecutionProvider laufen, werden span und Skalar als Eingaben unterstützt. Die Verwendung finden Sie hier.
  • Für benutzerdefinierte Operatorfunktionen, die einen Kernel-Kontext erwarten, siehe ein Beispiel hier.
  • Wenn Sie unique_ptr verwenden, um einen erstellten benutzerdefinierten Operator zu hosten, stellen Sie sicher, dass dieser zusammen mit der verbrauchenden Sitzung erhalten bleibt.

Weitere Beispiele finden Sie hier und hier.

Legacy-Methode für die Entwicklung und Registrierung benutzerdefinierter Operatoren

Die Legacy-Methode für die Entwicklung benutzerdefinierter Operatoren wird weiterhin unterstützt. Bitte beachten Sie die Beispiele hier.

Erstellen einer Bibliothek benutzerdefinierter Operatoren

Benutzerdefinierte Operatoren können in einer separaten freigegebenen Bibliothek (z. B. einer .dll unter Windows oder einer .so unter Linux) definiert werden. Eine benutzerdefinierte Operatorbibliothek muss eine Funktion namens RegisterCustomOps exportieren und implementieren. Die Funktion RegisterCustomOps fügt eine Ort::Custom::OpDomain hinzu, die die benutzerdefinierten Operatoren der Bibliothek enthält, zu den bereitgestellten Sitzungsoptionen. Bitte beachten Sie ein Projekt hier und zugehörige CMake-Befehle hier.

Aufrufen eines nativen Operators aus einem benutzerdefinierten Operator

Um die Implementierung benutzerdefinierter Operatoren zu vereinfachen, können native ONNX Runtime-Operatoren direkt aufgerufen werden. Einige benutzerdefinierte Operatoren müssen möglicherweise GEMM oder TopK zwischen anderen Berechnungen durchführen. Dies kann auch für die Vorverarbeitung und Nachbearbeitung eines Knotens wie Conv für die Zustandsverwaltung nützlich sein. Um dies zu erreichen, kann der Conv-Knoten von einem benutzerdefinierten Operator wie CustomConv umschlossen werden, innerhalb dessen die Eingabe und Ausgabe zwischengespeichert und verarbeitet werden können.

Diese Funktion wird ab ONNX Runtime 1.12.0+ unterstützt. Siehe: API und Beispiele.

Benutzerdefinierte Ops für CUDA und ROCM

Seit onnxruntime 1.16 werden kundenindividuelle Ops für CUDA- und ROCM-Geräte unterstützt. Gerätebezogene Ressourcen können direkt innerhalb des Ops über einen gerätebezogenen Kontext abgerufen werden. Nehmen wir CUDA als Beispiel

void KernelOne(const Ort::Custom::CudaContext& cuda_ctx,
               const Ort::Custom::Tensor<float>& X,
               const Ort::Custom::Tensor<float>& Y,
               Ort::Custom::Tensor<float>& Z) {
  auto input_shape = X.Shape();
  CUSTOM_ENFORCE(cuda_ctx.cuda_stream, "failed to fetch cuda stream");
  CUSTOM_ENFORCE(cuda_ctx.cudnn_handle, "failed to fetch cudnn handle");
  CUSTOM_ENFORCE(cuda_ctx.cublas_handle, "failed to fetch cublas handle");
  auto z_raw = Z.Allocate(input_shape);
  cuda_add(Z.NumberOfElement(), z_raw, X.Data(), Y.Data(), cuda_ctx.cuda_stream); // launch a kernel inside
}

Ein vollständiges Beispiel finden Sie hier. Um die Entwicklung weiter zu erleichtern, werden eine Vielzahl von CUDA EP-Ressourcen und Konfigurationen über CudaContext bereitgestellt. Einzelheiten finden Sie im Header.

Für ROCM ist es ähnlich wie

void KernelOne(const Ort::Custom::RocmContext& rocm_ctx,
               const Ort::Custom::Tensor<float>& X,
               const Ort::Custom::Tensor<float>& Y,
               Ort::Custom::Tensor<float>& Z) {
  auto input_shape = X.Shape();
  CUSTOM_ENFORCE(rocm_ctx.hip_stream, "failed to fetch hip stream");
  CUSTOM_ENFORCE(rocm_ctx.miopen_handle, "failed to fetch miopen handle");
  CUSTOM_ENFORCE(rocm_ctx.rblas_handle, "failed to fetch rocblas handle");
  auto z_raw = Z.Allocate(input_shape);
  rocm_add(Z.NumberOfElement(), z_raw, X.Data(), Y.Data(), rocm_ctx.hip_stream); // launch a kernel inside
}

Details finden Sie hier.

Ein Operator, verschiedene Typen

Seit onnxruntime 1.16 ist es einem benutzerdefinierten Operator gestattet, verschiedene Datentypen zu unterstützen

template <typename T>
void MulTop(const Ort::Custom::Span<T>& in, Ort::Custom::Tensor<T>& out) {
  out.Allocate({1})[0] = in[0] * in[1];
}

int main() {
  std::unique_ptr<OrtLiteCustomOp> c_MulTopOpFloat{Ort::Custom::CreateLiteCustomOp("MulTop", "CPUExecutionProvider", MulTop<float>)};
  std::unique_ptr<OrtLiteCustomOp> c_MulTopOpInt32{Ort::Custom::CreateLiteCustomOp("MulTop", "CPUExecutionProvider", MulTop<int32_t>)};
  // create a domain adding both c_MulTopOpFloat and c_MulTopOpInt32
}

Code finden Sie hier. Ein Unit-Testfall finden Sie hier.

Einbinden einer externen Inferenz-Runtime in einen benutzerdefinierten Operator

Ein benutzerdefinierter Operator kann ein ganzes Modell umschließen, das dann mit einer externen API oder Runtime inferenziert wird. Dies kann die Integration externer Inferenz-Engines oder APIs mit ONNX Runtime erleichtern.

Betrachten Sie als Beispiel das folgende ONNX-Modell mit einem benutzerdefinierten Operator namens „OpenVINO_Wrapper“. Der „OpenVINO_Wrapper“-Knoten kapselt ein ganzes MNIST-Modell im nativen Modellformat von OpenVINO (XML- und BIN-Daten). Die Modelldaten werden in die Attribute des Knotens serialisiert und später vom Kernel des benutzerdefinierten Operators abgerufen, um eine In-Memory-Darstellung des Modells zu erstellen und die Inferenz mit OpenVINO C++ APIs auszuführen.

ONNX model of a custom operator wrapping an OpenVINO MNIST model

Der folgende Codeausschnitt zeigt, wie der benutzerdefinierte Operator definiert ist.

// Note - below code utilizes legacy custom op interfaces
struct CustomOpOpenVINO : Ort::CustomOpBase<CustomOpOpenVINO, KernelOpenVINO> {
  explicit CustomOpOpenVINO(Ort::ConstSessionOptions session_options);

  CustomOpOpenVINO(const CustomOpOpenVINO&) = delete;
  CustomOpOpenVINO& operator=(const CustomOpOpenVINO&) = delete;

  void* CreateKernel(const OrtApi& api, const OrtKernelInfo* info) const;

  constexpr const char* GetName() const noexcept {
    return "OpenVINO_Wrapper";
  }

  constexpr const char* GetExecutionProviderType() const noexcept {
    return "CPUExecutionProvider";
  }

  // IMPORTANT: In order to wrap a generic runtime-specific model, the custom operator
  // must have a single non-homogeneous variadic input and output.

  constexpr size_t GetInputTypeCount() const noexcept {
    return 1;
  }

  constexpr size_t GetOutputTypeCount() const noexcept {
    return 1;
  }

  constexpr ONNXTensorElementDataType GetInputType(size_t /* index */) const noexcept {
    return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
  }

  constexpr ONNXTensorElementDataType GetOutputType(size_t /* index */) const noexcept {
    return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
  }

  constexpr OrtCustomOpInputOutputCharacteristic GetInputCharacteristic(size_t /* index */) const noexcept {
    return INPUT_OUTPUT_VARIADIC;
  }

  constexpr OrtCustomOpInputOutputCharacteristic GetOutputCharacteristic(size_t /* index */) const noexcept {
    return INPUT_OUTPUT_VARIADIC;
  }

  constexpr bool GetVariadicInputHomogeneity() const noexcept {
    return false;  // heterogenous
  }

  constexpr bool GetVariadicOutputHomogeneity() const noexcept {
    return false;  // heterogeneous
  }

  // The "device_type" is configurable at the session level.
  std::vector<std::string> GetSessionConfigKeys() const { return {"device_type"}; }

 private:
  std::unordered_map<std::string, std::string> session_configs_;
};

Beachten Sie, dass der benutzerdefinierte Operator so definiert ist, dass er eine einzelne variadische/heterogene Eingabe und eine einzelne variadische/heterogene Ausgabe hat. Dies ist notwendig, um OpenVINO-Modelle mit unterschiedlichen Eingabe- und Ausgabetypen und -formen (nicht nur ein MNIST-Modell) umschließen zu können. Weitere Informationen zu Ein- und Ausgabeeigenschaften finden Sie in der Dokumentation zur OrtCustomOp-Struktur.

Zusätzlich deklariert der benutzerdefinierte Operator „device_type“ als Sitzungskonfiguration, die von der Anwendung festgelegt werden kann. Der folgende Codeausschnitt zeigt, wie eine benutzerdefinierte Operatorbibliothek, die den oben genannten benutzerdefinierten Operator enthält, registriert und konfiguriert wird.

Ort::Env env;
Ort::SessionOptions session_options;
Ort::CustomOpConfigs custom_op_configs;

// Create local session config entries for the custom op.
custom_op_configs.AddConfig("OpenVINO_Wrapper", "device_type", "CPU");

// Register custom op library and pass in the custom op configs (optional).
session_options.RegisterCustomOpsLibrary("MyOpenVINOWrapper_Lib.so", custom_op_configs);

Ort::Session session(env, ORT_TSTR("custom_op_mnist_ov_wrapper.onnx"), session_options);

Weitere Details finden Sie im vollständigen Beispiel für den OpenVINO Custom Operator Wrapper. Um ein ONNX-Modell zu erstellen, das ein externes Modell oder Gewichte umschließt, beachten Sie das Tool create_custom_op_wrapper.py.