Benutzerdefinierte Operatoren mit Python-Funktionen erstellen

Benutzerdefinierte Operatoren sind ein mächtiges Merkmal von ONNX Runtime, das es Benutzern ermöglicht, die Funktionalität der Laufzeitumgebung zu erweitern, indem sie eigene Operatoren implementieren, um spezifische Operationen durchzuführen, die im Standard-ONNX-Operatorkatalog nicht verfügbar sind.

In diesem Dokument stellen wir vor, wie man einen benutzerdefinierten Operator mit Python-Funktionen erstellt und ihn zur Inferenz in ONNX Runtime integriert.

Schritt 1: Definieren Sie die Python-Funktion für den benutzerdefinierten Operator

Beginnen Sie mit der Definition der Python-Funktion, die als Implementierung für Ihren benutzerdefinierten Operator dienen wird. Stellen Sie sicher, dass die Funktion mit den Eingabe- und Ausgabetensorformen kompatibel ist, die Sie für Ihren benutzerdefinierten Operator erwarten. Der Python-Decorator @onnx_op konvertiert die Funktion in eine Implementierung für einen benutzerdefinierten Operator. Das Folgende ist ein Beispiel, bei dem wir eine Funktion für einen Tokenizer erstellen.

@onnx_op(op_type="GPT2Tokenizer",
            inputs=[PyCustomOpDef.dt_string],
            outputs=[PyCustomOpDef.dt_int64, PyCustomOpDef.dt_int64],
            attrs={"padding_length": PyCustomOpDef.dt_int64})
def bpe_tokenizer(s, **kwargs):
    padding_length = kwargs["padding_length"]
    input_ids, attention_mask = cls.tokenizer.tokenizer_sentence([s[0]], padding_length)
    return input_ids, attention_mask

Da ONNXRuntimme das benutzerdefinierte Operatorenschema beim Laden eines Modells benötigt, geben Sie diese bitte über die Argumente von onnx_op an. Außerdem ist 'attrs' erforderlich, wenn es Attribute für den ONNX-Knoten gibt, die ein Wörterbuch sein können, das vom Namen zum Typ zugeordnet ist, oder eine Liste, wenn alle Typen nur Zeichenketten sind.

Schritt 2: Erstellen Sie ein ONNX-Modell mit dem benutzerdefinierten Operator

Nachdem der benutzerdefinierte Operator bei ONNX Runtime registriert ist, können Sie ein ONNX-Modell erstellen, das ihn verwendet. Sie können entweder ein bestehendes ONNX-Modell modifizieren, um den benutzerdefinierten Operator einzubinden, oder ein neues von Grund auf neu erstellen.

Um ein neues ONNX-Modell mit dem benutzerdefinierten Operator zu erstellen, können Sie die ONNX Python API verwenden. Hier ist ein Beispiel: test_pyops.py

Einen benutzerdefinierten Operator von Grund auf in C++ erstellen

Bevor Sie einen benutzerdefinierten Operator implementieren, benötigen Sie ein ONNX-Modell mit einem oder mehreren ORT-benutzerdefinierten Operatoren, das von ONNX-Konvertern wie ONNX-Script, ONNX Model API usw. erstellt wurde.

1. Schnelle Überprüfung mit PythonOp (optional)

Bevor Sie tatsächlich einen benutzerdefinierten Operator für Ihren Anwendungsfall entwickeln, können Sie, wenn Sie das ONNX-Modell schnell mit Python überprüfen möchten, den benutzerdefinierten Operator mit Python-Funktionen wie oben beschrieben wrappen.

import numpy
from onnxruntime_extensions import PyOp, onnx_op

# Implement the CustomOp by decorating a function with onnx_op
@onnx_op(op_type="Inverse", inputs=[PyOp.dt_float])
def inverse(x):
    # the user custom op implementation here:
    return numpy.linalg.inv(x)

# Run the model with this custom op
# model_func = PyOrtFunction(model_path)
# outputs = model_func(inputs)
# ...

2. Generieren Sie den C++-Vorlagencode des benutzerdefinierten Operators aus dem ONNX-Modell (optional)

python -m onnxruntime-extensions.cmd --cpp-gen <model_path> <repository_dir>` If you are familiar with the ONNX model detail, you create the custom operator C++ classes directly.

3. Implementieren Sie die CustomOp Kernel Compute-Methode in den generierten C++-Dateien.

Der C++-Code für den benutzerdefinierten Operator-Kernel kann im Ordner operators gefunden werden, z. B. gaussian_blur. Alle C++-APIs, die in der Kernel-Implementierung verwendet werden können, sind unten aufgelistet.

  • ONNXRuntime Custom API-Dokumentation
  • Die in ONNXRuntime Extensions integrierten API-Dokumente von Drittanbieterbibliotheken, die im C++-Code verwendet werden können
    • OpenCV API-Dokumentation https://docs.opencv.org/4.x/
    • Google SentencePiece Library-Dokumentation https://github.com/google/sentencepiece/blob/master/doc/api.md
    • dlib (Matrix- und ML-Bibliothek) C++ API-Dokumentation https://dlib.net/algorithms.html
    • BlingFire Library https://github.com/microsoft/BlingFire
    • Google RE2 Library https://github.com/google/re2/wiki/CplusplusAPI
    • JSON-Bibliothek https://json.nlohmann.me/api/basic_json/

3. Erstellen und Testen

  • Die Unit-Tests können als Python oder C++ implementiert werden. Überprüfen Sie den Ordner test für weitere Beispiele.
  • Überprüfen Sie build-package, um zu erfahren, wie Sie das Sprachpaket für die Produktion erstellen.

Bitte überprüfen Sie die Beiträge, um zu sehen, ob es möglich ist, den benutzerdefinierten Operator zu onnxruntime-extensions beizutragen.