Benutzerdefinierte Operatoren

Da die integrierte Operatorbibliothek von TensorFlow Lite nur eine begrenzte Anzahl von TensorFlow-Operatoren unterstützt, ist nicht jedes Modell konvertierbar. Weitere Einzelheiten finden Sie unter Operator Kompatibilität .

Um die Konvertierung zu ermöglichen, können Benutzer ihre eigene benutzerdefinierte Implementierung eines nicht unterstützten TensorFlow-Operators in TensorFlow Lite bereitstellen, der als benutzerdefinierter Operator bezeichnet wird. Wenn stattdessen Sie wollen eine Reihe von nicht unterstützten (oder unterstützt) kombinieren TensorFlow Betreibern zu einem einzigen verschmolzen optimierte benutzerdefinierten Operator, siehe Operator Verschmelzen .

Die Verwendung benutzerdefinierter Operatoren besteht aus vier Schritten.

Lassen Sie sich zu Fuß durch ein End-to-End - Beispiel für ein Modell mit einem benutzerdefinierten Operator läuft tf.sin (benannt als Sin siehe #create_a_tensorflow_model) , die in TensorFlow unterstützt wird, aber nicht unterstütze in TensorFlow Lite.

Beispiel: Benutzerdefinierte Sin Operator

Sehen wir uns ein Beispiel für die Unterstützung eines TensorFlow-Operators an, den TensorFlow Lite nicht hat. Angenommen , wir das verwenden Sin Operator und dass wir bauen ein sehr einfaches Modell für eine Funktion y = sin(x + offset) , in dem offset ist trainierbar.

Erstellen Sie ein TensorFlow-Modell

Der folgende Codeausschnitt trainiert ein einfaches TensorFlow-Modell. Dieses Modell enthält nur einen benutzerdefinierten Operator namens Sin , die eine Funktion ist y = sin(x + offset) , in der offset - trainierbar ist.

import tensorflow as tf

# Define training dataset and variables
x = [-8, 0.5, 2, 2.2, 201]
y = [-0.6569866 ,  0.99749499,  0.14112001, -0.05837414,  0.80641841]
offset = tf.Variable(0.0)

# Define a simple model which just contains a custom operator named `Sin`
@tf.function
def sin(x):
  return tf.sin(x + offset, name="Sin")

# Train model
optimizer = tf.optimizers.Adam(0.01)
def train(x, y):
    with tf.GradientTape() as t:
      predicted_y = sin(x)
      loss = tf.reduce_sum(tf.square(predicted_y - y))
    grads = t.gradient(loss, [offset])
    optimizer.apply_gradients(zip(grads, [offset]))

for i in range(1000):
    train(x, y)

print("The actual offset is: 1.0")
print("The predicted offset is:", offset.numpy())
The actual offset is: 1.0
The predicted offset is: 1.0000001

Wenn Sie an dieser Stelle versuchen, ein TensorFlow Lite-Modell mit den Standardkonverter-Flags zu generieren, erhalten Sie die folgende Fehlermeldung:

Error:
Some of the operators in the model are not supported by the standard TensorFlow
Lite runtime...... Here is
a list of operators for which you will need custom implementations: Sin.

In ein TensorFlow Lite-Modell umwandeln

Erstellen Sie ein TensorFlow Lite - Modell mit benutzerdefinierten Operatoren, indem Sie die Konverter - Attribut Einstellung allow_custom_ops wie unten dargestellt:

converter = tf.lite.TFLiteConverter.from_concrete_functions([sin.get_concrete_function(x)], sin)
converter.allow_custom_ops = True
tflite_model = converter.convert()

Wenn Sie es jetzt mit dem Standardinterpreter ausführen, erhalten Sie die folgenden Fehlermeldungen:

Error:
Didn't find custom operator for name 'Sin'
Registration failed.

Erstellen und registrieren Sie den Operator.

Alle TensorFlow Lite-Operatoren (sowohl benutzerdefinierte als auch integrierte) werden über eine einfache reine C-Schnittstelle definiert, die aus vier Funktionen besteht:

typedef struct {
  void* (*init)(TfLiteContext* context, const char* buffer, size_t length);
  void (*free)(TfLiteContext* context, void* buffer);
  TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node);
  TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node);
} TfLiteRegistration;

Siehe common.h Einzelheiten über TfLiteContext und TfLiteNode . Ersteres bietet Möglichkeiten zur Fehlerberichterstattung und Zugriff auf globale Objekte, einschließlich aller Tensoren. Letzteres ermöglicht Implementierungen den Zugriff auf ihre Ein- und Ausgänge.

Wenn der Interpreter lädt ein Modell, ruft es init() einmal für jeden Knoten im Graphen. Eine gegebene init() wird mehr als einmal aufgerufen , wenn der op mehrfach in der Grafik verwendet wird. Für benutzerdefinierte Operationen wird ein Konfigurationspuffer bereitgestellt, der einen Flexbuffer enthält, der Parameternamen ihren Werten zuordnet. Der Puffer ist für eingebaute Ops leer, da der Interpreter die Op-Parameter bereits geparst hat. Kernel-Implementierungen, die einen Status erfordern, sollten ihn hier initialisieren und den Besitz an den Aufrufer übertragen. Für jede init() Aufruf, gibt es einen entsprechenden Anruf sein free() , so dass Implementierungen der Puffer verfügen sie in zugewiesen haben könnte init() .

Immer wenn die Größe der Eingabetensoren geändert wird, durchläuft der Interpreter den Graphen und benachrichtigt die Implementierungen über die Änderung. Dies gibt ihnen die Möglichkeit, die Größe ihres internen Puffers zu ändern, die Gültigkeit von Eingabe-Shapes und -Typen zu überprüfen und Ausgabe-Shapes neu zu berechnen. Dies alles wird durch getan prepare() , und Implementierungen können ihren Zustand Zugriff mit node->user_data .

Schließlich wird jedes Mal Inferenz ausgeführt wird , durchläuft der Interpreter die Grafik calling invoke() , und auch hier ist der Staat als node->user_data .

Benutzerdefinierte Ops können genauso wie eingebaute Ops implementiert werden, indem diese vier Funktionen und eine globale Registrierungsfunktion definiert werden, die normalerweise wie folgt aussieht:

namespace tflite {
namespace ops {
namespace custom {
  TfLiteRegistration* Register_MY_CUSTOM_OP() {
    static TfLiteRegistration r = {my_custom_op::Init,
                                   my_custom_op::Free,
                                   my_custom_op::Prepare,
                                   my_custom_op::Eval};
    return &r;
  }
}  // namespace custom
}  // namespace ops
}  // namespace tflite

Beachten Sie, dass die Registrierung nicht automatisch und ein expliziter Aufruf ist Register_MY_CUSTOM_OP gemacht werden soll. Während die Standard - BuiltinOpResolver (erhältlich bei der :builtin_ops Ziel) Pflege der Registrierung von builtins nimmt, werden individuelle ops haben in separaten benutzerdefinierten Bibliotheken gesammelt werden.

Definieren des Kernels in der TensorFlow Lite-Laufzeit

Alles , was wir brauchen die op in TensorFlow Lite zu verwenden , zu tun ist , zwei Funktionen definieren ( Prepare und Eval ) und ein Konstrukt TfLiteRegistration :

TfLiteStatus SinPrepare(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

  const TfLiteTensor* input = GetInput(context, node, 0);
  TfLiteTensor* output = GetOutput(context, node, 0);

  int num_dims = NumDimensions(input);

  TfLiteIntArray* output_size = TfLiteIntArrayCreate(num_dims);
  for (int i=0; i<num_dims; ++i) {
    output_size->data[i] = input->dims->data[i];
  }

  return context->ResizeTensor(context, output, output_size);
}

TfLiteStatus SinEval(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  const TfLiteTensor* input = GetInput(context, node,0);
  TfLiteTensor* output = GetOutput(context, node,0);

  float* input_data = input->data.f;
  float* output_data = output->data.f;

  size_t count = 1;
  int num_dims = NumDimensions(input);
  for (int i = 0; i < num_dims; ++i) {
    count *= input->dims->data[i];
  }

  for (size_t i=0; i<count; ++i) {
    output_data[i] = sin(input_data[i]);
  }
  return kTfLiteOk;
}

TfLiteRegistration* Register_SIN() {
  static TfLiteRegistration r = {nullptr, nullptr, SinPrepare, SinEval};
  return &r;
}

Wenn die Initialisierung OpResolver , fügen Sie den benutzerdefinierten op in den Resolver (Beispiel unten sehen). Dadurch wird der Operator bei Tensorflow Lite registriert, damit TensorFlow Lite die neue Implementierung verwenden kann. Beachten Sie, dass die letzten beiden Argumente in TfLiteRegistration entspricht das SinPrepare und SinEval Funktionen , die Sie für die benutzerdefinierte op definiert. Wenn Sie verwendet SinInit und SinFree Funktionen Variablen im OP verwendet zu initialisieren und um Speicherplatz freizugeben, bzw. dann würden sie auf die ersten beiden Argumente hinzugefügt werden TfLiteRegistration ; Diese Argumente sind zu setzen nullptr in diesem Beispiel.

Registrieren Sie den Operator bei der Kernel-Bibliothek

Jetzt müssen wir den Operator bei der Kernel-Bibliothek registrieren. Dies wird mit einem getan OpResolver . Hinter den Kulissen lädt der Interpreter eine Bibliothek von Kerneln, die zugewiesen werden, um jeden der Operatoren im Modell auszuführen. Während die Standardbibliothek nur eingebaute Kernel enthält, ist es möglich, sie durch eine benutzerdefinierte Bibliothek op-Operatoren zu ersetzen/zu erweitern.

Die OpResolver - Klasse, die Operatorcodes und Namen in tatsächlichen Code übersetzt wird wie folgt definiert:

class OpResolver {
  virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
  virtual TfLiteRegistration* FindOp(const char* op) const = 0;
  virtual void AddBuiltin(tflite::BuiltinOperator op, TfLiteRegistration* registration) = 0;
  virtual void AddCustom(const char* op, TfLiteRegistration* registration) = 0;
};

Regelmäßige Anwendung erfordert , dass Sie die Verwendung BuiltinOpResolver und schreiben:

tflite::ops::builtin::BuiltinOpResolver resolver;

So fügen Sie die benutzerdefinierte op oben erstellt, rufen Sie AddOp (bevor Sie den Resolver an den Pass InterpreterBuilder ):

resolver.AddCustom("Sin", Register_SIN());

Wenn die Menge des eingebauten ops gilt als zu groß sein, eine neue OpResolver könnte Code erzeugt wird basierend auf einer bestimmte Untergruppe von ops, möglicherweise nur die , die in einem bestimmten Modell enthalten ist . Dies ist das Äquivalent von TensorFlow selektiver Registrierung (und eine einfache Version davon ist in der zur Verfügung stehenden tools - Verzeichnis).

Wenn Sie Ihre benutzerdefinierten Operatoren in Java definieren wollen, würden Sie zur Zeit benötigen , um Ihre eigenen benutzerdefinierten JNI - Schicht zu bauen und eigene AAR kompilieren in diesem jni Code . Und falls Sie diese Operatoren in Python definieren möchten , können Sie Ihre Eintragungen im platzieren Python - Wrapper - Code .

Beachten Sie, dass ein ähnlicher Prozess wie oben befolgt werden kann, um einen Satz von Operationen anstelle eines einzelnen Operators zu unterstützen. Fügen Sie einfach so viele AddCustom Operatoren wie Sie benötigen. Darüber hinaus BuiltinOpResolver können auch mit dem Sie Implementierungen von builtins außer Kraft zu setzen AddBuiltin .

Testen und profilieren Sie Ihren Betreiber

Zum Profil Ihrer op mit dem TensorFlow Lite - Benchmark - Tool, können Sie das verwenden Benchmark - Modell Tool für TensorFlow Lite. Zu Testzwecken können Sie Ihre lokale Build von TensorFlow Lite bewusst Ihrer individuellen op machen , indem Sie den entsprechenden Zugabe AddCustom Aufrufs (als Erscheinen oben) register.cc

Empfohlene Vorgehensweise

  1. Optimieren Sie Speicherzuweisungen und -aufhebungen vorsichtig. Zuweisen von Speicher in Prepare ist effizienter als in Invoke , und vor einem Schleifenspeicher Zuteilen ist besser als in jeder Iteration. Verwenden Sie temporäre Tensordaten, anstatt sich selbst zu verfehlen (siehe Punkt 2). Verwenden Sie Zeiger/Referenzen, anstatt so viel wie möglich zu kopieren.

  2. Wenn eine Datenstruktur während des gesamten Vorgangs bestehen bleibt, empfehlen wir, den Speicher mit temporären Tensoren vorab zuzuweisen. Möglicherweise müssen Sie die OpData-Struktur verwenden, um auf die Tensorindizes in anderen Funktionen zu verweisen. Siehe das Beispiel in der Kernel für Faltung . Ein Beispiel-Code-Snippet finden Sie unten

    auto* op_data = reinterpret_cast<OpData*>(node->user_data);
    TfLiteIntArrayFree(node->temporaries);
    node->temporaries = TfLiteIntArrayCreate(1);
    node->temporaries->data[0] = op_data->temp_tensor_index;
    TfLiteTensor* temp_tensor = &context->tensors[op_data->temp_tensor_index];
    temp_tensor->type =  kTfLiteFloat32;
    temp_tensor->allocation_type = kTfLiteArenaRw;
    
  3. Wenn es nicht zu viel kostet verschwendet Speicher, lieber eine statische feste Größe Array (oder eine vorab zugewiesene std::vector in Resize ) , anstatt einen dynamisch zugewiesenen mit std::vector jeder Iteration der Ausführung.

  4. Vermeiden Sie die Instanziierung von Standardbibliothekscontainervorlagen, die noch nicht vorhanden sind, da sie die Binärgröße beeinflussen. Zum Beispiel, wenn Sie ein brauchen std::map in Ihrem Betrieb , die in anderen Kernel nicht vorhanden ist , unter Verwendung eine std::vector mit direkter Indizierung Mapping könnte arbeiten , während die binäre Größe klein zu halten. Sehen Sie, was andere Kernel verwenden, um Erkenntnisse zu gewinnen (oder fragen Sie).

  5. Überprüfen Sie den Zeiger auf die von zurückgeführtem Speicher malloc . Wenn dieser Zeiger ist nullptr sollten keine Operationen durchgeführt werden , dass die Zeiger verwenden. Wenn Sie malloc in einer Funktion und einen Fehlerausgang haben, ausplanen Speicher , bevor Sie verlassen.

  6. Verwenden TF_LITE_ENSURE(context, condition) für einen bestimmten Zustand zu überprüfen. Ihr Code muss nicht Gedächtnis hängen lassen , wenn TF_LITE_ENSURE verwendet wird, das heißt, diese Makros verwendet werden sollten , bevor Ressourcen zugewiesen werden , die auslaufen wird.