Weitere Informationen zu den Ankündigungen der Google AI Edge auf der Google I/O

Diese Seite wurde von der Cloud Translation API übersetzt.

Benutzerdefinierte Operatoren

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

Um eine Konvertierung zu ermöglichen, können Nutzer eine eigene benutzerdefinierte Implementierung eines nicht unterstützten TensorFlow-Operators in TensorFlow Lite bereitstellen, der als benutzerdefinierter Operator bezeichnet wird. Wenn Sie stattdessen eine Reihe nicht unterstützter (oder unterstützter) TensorFlow-Operatoren zu einem einzelnen kombinierten optimierten benutzerdefinierten Operator kombinieren möchten, finden Sie weitere Informationen unter Operator „Fusing“.

Benutzerdefinierte Operatoren werden in vier Schritten verwendet.

Erstellen Sie ein TensorFlow-Modell. Achten Sie darauf, dass das gespeicherte Modell (oder Graph Def) auf den korrekt benannten TensorFlow Lite-Operator verweist.
In ein TensorFlow Lite-Modell konvertieren Achten Sie darauf, dass Sie das richtige TensorFlow Lite-Converter-Attribut festgelegt haben, damit das Modell erfolgreich konvertiert werden kann.
Erstellen und registrieren Sie den Operator. So weiß die TensorFlow Lite-Laufzeit, wie die Operatoren und Parameter in Ihrer Grafik ausführbarem C/C++-Code zugeordnet werden.
Anbieter testen und Profil erstellen Wenn Sie nur den benutzerdefinierten Operator testen möchten, erstellen Sie am besten ein Modell nur mit diesem benutzerdefinierten Operator und verwenden das Programm benchmark_model.

Sehen wir uns ein End-to-End-Beispiel für die Ausführung eines Modells mit dem benutzerdefinierten Operator tf.atan (namens Atan, siehe TensorFlow-Modell erstellen) an, der in TensorFlow unterstützt wird, in TensorFlow Lite aber nicht.

Der TensorFlow-Operator „Text“ ist ein Beispiel für einen benutzerdefinierten Operator. Ein Codebeispiel finden Sie in der Anleitung TF-Text in TF Lite konvertieren.

Beispiel: Benutzerdefinierter `Atan`-Operator

Sehen wir uns ein Beispiel für die Unterstützung eines TensorFlow-Operators an, den TensorFlow Lite nicht hat. Angenommen, wir verwenden den Operator Atan und erstellen ein sehr einfaches Modell für die Funktion y = atan(x + offset), in der offset trainierbar ist.

TensorFlow-Modell erstellen

Mit dem folgenden Code-Snippet wird ein einfaches TensorFlow-Modell trainiert. Dieses Modell enthält nur einen benutzerdefinierten Operator namens Atan, bei dem es sich um eine Funktion y = atan(x + offset) handelt, bei der offset trainierbar ist.

import tensorflow as tf

# Define training dataset and variables
x = [-8, 0.5, 2, 2.2, 201]
y = [-1.4288993, 0.98279375, 1.2490457, 1.2679114, 1.5658458]
offset = tf.Variable(0.0)

# Define a simple model which just contains a custom operator named `Atan`
@tf.function(input_signature=[tf.TensorSpec.from_tensor(tf.constant(x))])
def atan(x):
  return tf.atan(x + offset, name="Atan")

# Train model
optimizer = tf.optimizers.Adam(0.01)
def train(x, y):
    with tf.GradientTape() as t:
      predicted_y = atan(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: 0.99999905

Wenn Sie an dieser Stelle versuchen, ein TensorFlow Lite-Modell mit den Standard-Converter-Flags zu generieren, wird die folgende Fehlermeldung angezeigt:

Error:
error: 'tf.Atan' op is neither a custom op nor a flex op.

In ein TensorFlow Lite-Modell konvertieren

Erstellen Sie ein TensorFlow Lite-Modell mit benutzerdefinierten Operatoren. Legen Sie dazu das Konvertierungsattribut allow_custom_ops wie unten gezeigt fest:

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

Wenn Sie es jetzt mit dem Standardinterpreter mit Befehlen wie den folgenden ausführen:

interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

Der Fehler wird weiterhin angezeigt:

Encountered unresolved custom op: Atan.

Erstellen und registrieren Sie den Operator.

#include "third_party/tensorflow/lite/c/c_api.h"
#include "third_party/tensorflow/lite/c/c_api_opaque.h"

Benutzerdefinierte TensorFlow Lite-Operatoren werden mithilfe einer einfachen reinen C-API definiert, die aus einem intransparenten Typ (TfLiteRegistrationExternal) und zugehörigen Funktionen besteht.

TfLiteRegistrationExternal ist ein intransparenter Typ:

typedef struct TfLiteRegistrationExternal TfLiteRegistrationExternal;

TfLiteRegistrationExternal speichert die Identität und Implementierung des Operators. Der Operator unterscheidet sich von seinen Operanden, die in den TF Lite-Grafikknoten für Knoten gespeichert werden, die den Operator aufrufen.

Instanzen dieses Typs werden mit Aufrufen von TfLiteRegistrationExternalCreate erstellt und können durch Aufrufen von TfLiteRegistrationExternalDelete gelöscht werden.

Die Identität des Operators wird über die Parameter für die Konstruktorfunktion TfLiteRegistrationExternalCreate festgelegt:

TfLiteRegistrationExternal*
TfLiteRegistrationExternalCreate(
    TfLiteBuiltinOperator builtin_code,  // Normally `TfLiteBuiltinCustom`.
    const char* custom_name,  // The name of the custom op.
    int version  // Normally `1` for the first version of a custom op.
);

Bei der Operatorimplementierung können „Methoden“ mit den folgenden Signaturen definiert werden. Alle diese Methoden sind optional. Damit ein Operator jedoch erfolgreich ausgewertet werden kann, muss die Operatorimplementierung mindestens die Methoden Prepare und Invoke definieren und festlegen (mithilfe der Setter-Funktionen).

// Initializes the op from serialized data.
void* Init(TfLiteOpaqueContext* context, const char* buffer, size_t length);

// Deallocates the op.
// The pointer `buffer` is the data previously returned by an Init invocation.
void Free(TfLiteOpaqueContext* context, void* buffer);

// Called when the inputs that this node depends on have been resized.
TfLiteStatus Prepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node);

// Called when the node is executed. (Should read node inputs and write to
// node outputs).
TfLiteStatus Invoke(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node);

// Retrieves the async kernel.
TfLiteAsyncKernel AsyncKernel(TfLiteOpaqueContext* context,
                              TfLiteOpaqueNode* node);

Die names (oder Namespace-Präfixe bei C++) in Ihrer Vorgangsimplementierung müssen nicht mit den Funktionsnamen im obigen Code-Snippet übereinstimmen, da die Custom Ops API für TF Lite nur ihre Adressen verwendet. Wir empfehlen, sie in einem anonymen Namespace oder als statische Funktionen zu deklarieren.

Es empfiehlt sich jedoch, den Operatornamen als Namespace oder Präfix in diese Funktionsnamen aufzunehmen:

C++

namespace my_namespace::my_custom_op {
  void* Init(TfLiteOpaqueContext* context,
             const char* buffer, size_t length) { ... }
  // ... plus definitions of Free, Prepare, and Invoke ...
}

C

void* MyCustomOpInit(TfLiteOpaqueContext* context,
                     const char* buffer, size_t length) { ... }
// ... plus definitions of MyCustomOpFree, MyCustomOpPrepare, and
// MyCustomOpInvoke.

Da es sich um eine C API handelt, werden diese „Methoden“ als C-Funktionszeiger im Typ TfLiteRegistrationExternal implementiert. Diese werden durch Übergeben der Adressen der Implementierungsfunktionen an die entsprechenden Setter-Funktionen TfLiteRegistrationExternalSetMethodName festgelegt:

void TfLiteRegistrationExternalSetInit(
    TfLiteRegistrationExternal* registration,
    void* (*init)(TfLiteOpaqueContext* context, const char* buffer,
                  size_t length));
void TfLiteRegistrationExternalSetFree(
    TfLiteRegistrationExternal* registration,
    void (*free)(TfLiteOpaqueContext* context, void* data));
void TfLiteRegistrationExternalSetPrepare(
    TfLiteRegistrationExternal* registration,
    TfLiteStatus (*prepare)(TfLiteOpaqueContext* context,
                            TfLiteOpaqueNode* node));
void TfLiteRegistrationExternalSetInvoke(
    TfLiteRegistrationExternal* registration,
    TfLiteStatus (*invoke)(TfLiteOpaqueContext* context,
                           TfLiteOpaqueNode* node));
void TfLiteRegistrationExternalSetAsyncKernel(
    TfLiteRegistrationExternal* registration,
    struct TfLiteAsyncKernel* (*async_kernel)(TfLiteOpaqueContext* context,
                                              TfLiteOpaqueNode* node));

Weitere Informationen zu TfLiteContext und TfLiteNode finden Sie unter common.h. TfLiteContext bietet Funktionen für Fehlerberichte und Zugriff auf globale Objekte, einschließlich aller Tensoren. TfLiteNode ermöglicht Operatorimplementierungen den Zugriff auf ihre Ein- und Ausgaben.

Wenn der Interpreter ein Modell lädt, ruft er die Methode Init() einmal für jeden Knoten im Diagramm auf. Eine bestimmte Init() wird mehr als einmal aufgerufen, wenn der Vorgang mehrmals im Diagramm verwendet wird. Für benutzerdefinierte Vorgänge wird ein Konfigurationspuffer bereitgestellt, der einen Flexbuffer enthält, der Parameternamen ihren Werten zuordnet. Der Zwischenspeicher ist für integrierte Vorgänge leer, da der Interpreter die Vorgangsparameter bereits geparst hat. Kernelimplementierungen, die einen Status erfordern, sollten ihn hier initialisieren und die Eigentümerschaft an den Aufrufer übertragen. Für jeden Init()-Aufruf gibt es einen entsprechenden Aufruf von Free(), sodass Implementierungen den Puffer entsorgen können, den sie möglicherweise in Init() zugewiesen haben.

Wenn die Größe der Eingabetensoren geändert wird, durchläuft der Interpreter die Grafik, die die Implementierungen über die Änderung informiert. Dadurch kann er die Größe des internen Puffers anpassen, die Gültigkeit von Eingabeformen und -typen prüfen und Ausgabeformen neu berechnen. Das erfolgt über die Methode Prepare() und Implementierungen können mit TfLiteOpaqueNodeGetUserData(node) auf ihren Status zugreifen.

Schließlich durchläuft der Interpreter bei jeder Inferenz die Grafik und ruft die Methode Invoke() auf. Auch hier ist der Status als TfLiteOpaqueNodeGetUserData(node) verfügbar.

Benutzerdefinierte Operationen können implementiert werden, indem diese „Methoden“-Funktionen und dann eine Funktion definiert wird, die eine Instanz von TfLiteRegistrationExternal zurückgibt, die durch Aufrufen von TfLiteRegistrationExternalCreate und der entsprechenden Setter-Methoden erstellt wurde:

C++

namespace my_namespace::my_custom_op {
  namespace {
    void* Init(TfLiteOpaqueContext* context,
               const char* buffer, size_t length) { ... }
    void Free(TfLiteOpaqueContext* context, void* buffer) { ... }
    TfLiteStatus Prepare(TfLiteOpaqueContext* context,
                         TfLiteOpaqueNode* node) { ... }
    TfLiteStatus Invoke(TfLiteOpaqueContext* context,
                        TfLiteOpaqueNode* node) {... }
  };

  const TfLiteRegistrationExternal* MyCustomOpRegistrationExternal() {
    // Singleton instance, intentionally never destroyed.
    static const TfLiteRegistrationExternal* my_custom_op = ()[] {
        TfLiteRegistrationExternal* r =
            TfLiteRegistrationExternalCreate(
                kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1);
        TfLiteRegistrationExternalSetInit(r, Init);
        TfLiteRegistrationExternalSetFree(r, Free);
        TfLiteRegistrationExternalSetPrepare(r, Prepare);
        TfLiteRegistrationExternalSetInvoke(r, Eval);
        return r;
      };
    return my_custom_op;
  }

  const TfLiteRegistration* MyCustomOpRegistration() {
    static const TfLiteRegistration my_custom_op {
      .registration_external = MyCustomOpRegistrationExternal();
    };
    return my_custom_op;
  }
}  // namespace my_namespace

C

static void* MyCustomOpInit(TfLiteOpaqueContext* context, const char* buffer,
                     size_t length) { ... }
static void MyCustomOpFree(TfLiteOpaqueContext* context, void* buffer) { ... }
static TfLiteStatus MyCustomOpPrepare(TfLiteOpaqueContext* context,
                                      TfLiteOpaqueNode* node) { ... }
static TfLiteStatus MyCustomOpInvoke(TfLiteOpaqueContext* context,
                                     TfLiteOpaqueNode* node) {... }

static TfLiteRegistrationExternal* MyCustomOpCreate() {
  const TfLiteRegistrationExternal* r =
      TfLiteRegistrationExternalCreate(
          kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1);
  TfLiteRegistrationExternalSetInit(r, MyCustomOpInit);
  TfLiteRegistrationExternalSetFree(r, MyCustomOpFree);
  TfLiteRegistrationExternalSetPrepare(r, MyCustomOpPrepare);
  TfLiteRegistrationExternalSetInvoke(r, MyCustomOpEval);
  return r;
}

const TfLiteRegistrationExternal* MyCustomOpRegistrationExternal() {
  // Singleton instance, intentionally never destroyed.
  static const TfLiteRegistrationExternal* my_custom_op = MyCustomOpCreate();
  return my_custom_op;
}

const TfLiteRegistration MyCustomOpRegistration() {
  static const TfLiteRegistration my_custom_op {
    .registration_external = MyCustomOpRegistrationExternal();
  };
  return my_custom_op;
}

Die Registrierung erfolgt nicht automatisch und die Funktion MyCustomOpRegistration sollte explizit aufgerufen werden (siehe Details unten). Während die Standard-BuiltinOpResolver (verfügbar über das Ziel :builtin_ops) die Registrierung von Builds übernimmt, müssen benutzerdefinierte Vorgänge in separaten benutzerdefinierten Bibliotheken erfasst werden.

Kernel in der TensorFlow Lite-Laufzeit definieren

Um den Vorgang in TensorFlow Lite zu verwenden, müssen Sie nur zwei Funktionen (Prepare und Eval) definieren und eine dritte nutzen, um eine TfLiteRegistrationExternal zu erstellen:

C++

namespace atan_op {
  namespace {
    TfLiteStatus AtanPrepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
      TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumInputs(node), 1);
      TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumOutputs(node), 1);

      const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
      TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

      int num_dims = TfLiteOpaqueTensorNumDimensions(input);

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

      return TfLiteOpaqueContextResizeTensor(context, output, output_size);
    }

    TfLiteStatus AtanEval(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
      const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
      TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

      float* input_data = static_cast(TfLiteOpaqueTensorData(input));
      float* output_data = static_cast(TfLiteOpaqueTensorData(output));

      size_t count = 1;
      int num_dims = TfLiteOpaqueTensorNumDimensions(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] = atan(input_data[i]);
      }
      return kTfLiteOk;
    }
  }  // anonymous namespace

  const TfLiteRegistrationExternal* AtanOpRegistrationExternal() {
    // Singleton instance, intentionally never destroyed.
    static const TfLiteRegistrationExternal* atan_op = ()[] {
        auto* r = TfLiteRegistrationExternalCreate(
            kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1);
        TfLiteRegistrationExternalSetPrepare(r, Prepare);
        TfLiteRegistrationExternalSetInvoke(r, Eval);
        return r;
      };
    return atan_op;
  }

  const TfLiteRegistration AtanOpRegistration() {
    static const TfLiteRegistration atan_op {
      .registration_external = AtanOpRegistrationExternal();
    };
    return atan_op;
  }
}  // namespace atan_op

C

static TfLiteStatus AtanPrepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
  TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumInputs(node), 1);
  TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumOutputs(node), 1);

  const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
  TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

  int num_dims = TfLiteOpaqueTensorNumDimensions(input);

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

  return TfLiteOpaqueContextResizeTensor(context, output, output_size);
}

static TfLiteStatus AtanEval(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
  const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
  TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

  float* input_data = static_cast(TfLiteOpaqueTensorData(input));
  float* output_data = static_cast(TfLiteOpaqueTensorData(output));

  size_t count = 1;
  int num_dims = TfLiteOpaqueTensorNumDimensions(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] = atan(input_data[i]);
  }
  return kTfLiteOk;
}

static const TfLiteRegistrationExternal* AtanOpCreate() {
  TfLiteRegistrationExternal* r = TfLiteRegistrationExternalCreate(
          kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1);
  TfLiteRegistrationExternalSetPrepare(r, Prepare);
  TfLiteRegistrationExternalSetInvoke(r, Eval);
  return r;
}

const TfLiteRegistrationExternal* AtanOpRegistrationExternal() {
  // Singleton instance, intentionally never destroyed.
  static const TfLiteRegistrationExternal* atan_op = AtanOpCreate();
  return atan_op;
}

const TfLiteRegistration AtanOpRegistration() {
  static const TfLiteRegistration atan_op {
    .registration_external = AtanOpRegistrationExternal();
  };
  return atan_op;
}

Fügen Sie beim Initialisieren von OpResolver die benutzerdefinierte Operation dem Resolver hinzu (siehe Beispiel unten). Dadurch wird der Operator bei TensorFlow Lite registriert, damit TensorFlow Lite die neue Implementierung verwenden kann. Die letzten beiden Argumente in TfLiteRegistration entsprechen den Funktionen AtanPrepare und AtanEval, die Sie für den benutzerdefinierten Vorgang definiert haben. Wenn Sie die Funktionen AtanInit und AtanFree verwendet haben, um im Vorgang verwendete Variablen zu initialisieren und Platz freizugeben, werden sie den ersten beiden Argumenten von TfLiteRegistration hinzugefügt. Diese Argumente sind in diesem Beispiel auf nullptr gesetzt.

Operator bei der Kernel-Bibliothek registrieren

Jetzt müssen wir den Operator bei der Kernel-Bibliothek registrieren. Dazu verwenden Sie OpResolver. Im Hintergrund lädt der Interpreter eine Bibliothek mit Kernels, die für die Ausführung der einzelnen Operatoren im Modell zugewiesen wird. Die Standardbibliothek enthält nur integrierte Kernel, kann aber durch Operatoren für benutzerdefinierte Bibliotheksvorgänge ersetzt oder erweitert werden.

Die Klasse OpResolver, die Operatorcodes und Namen in tatsächlichen Code umwandelt, ist so definiert:

class OpResolver {
 public:
  virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
  virtual TfLiteRegistration* FindOp(const char* op) const = 0;
  ...
};

Aus Gründen der Abwärtskompatibilität verwendet diese Klasse den älteren Betontyp TfLiteRegistration anstelle des intransparenten Typs TfLiteRegistrationExternal. Die Struktur TfLiteRegistration enthält jedoch ein registration_external-Feld vom Typ TfLiteRegistrationExternal*.

Die Klassen MutableOpResolver und BuiltinOpResolver werden von OpResolver abgeleitet:

class MutableOpResolver : public OpResolver {
 public:
  MutableOpResolver();  // Constructs an initially empty op resolver.
  void AddBuiltin(tflite::BuiltinOperator op, const TfLiteRegistration* registration) = 0;
  void AddCustom(const char* op, const TfLiteRegistration* registration) = 0;
  void AddAll(const MutableOpResolver& other);
  ...
};

class BuiltinOpResolver : public MutableOpResolver {
 public:
  BuiltinOpResolver();  // Constructs an op resolver with all the builtin ops.
};

Für die normale Nutzung (ohne benutzerdefinierte Vorgänge) müssen Sie den BuiltinOpResolver verwenden und schreiben:

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

Wenn Sie die oben erstellte benutzerdefinierte Operation hinzufügen möchten, können Sie stattdessen ein MutableOpResolver verwenden und AddCustom aufrufen (bevor Sie den Resolver an das InterpreterBuilder übergeben):

tflite::ops::builtin::MutableOpResolver resolver;
resolver.AddAll(tflite::ops::builtin::BuiltinOpResolver());
resolver.AddCustom("Atan", AtanOpRegistration());

Wenn die Menge der integrierten Operationen als zu groß erachtet wird, kann auf der Grundlage einer bestimmten Teilmenge von Vorgängen eine neue OpResolver generiert werden, möglicherweise nur anhand derjenigen, die in einem bestimmten Modell enthalten sind. Dies entspricht der selektiven Registrierung von TensorFlow. Eine einfache Version ist im Verzeichnis tools verfügbar.

Wenn Sie Ihre benutzerdefinierten Operatoren in Java definieren möchten, müssen Sie derzeit Ihre eigene benutzerdefinierte JNI-Ebene erstellen und in diesem JNI-Code Ihre eigene AAR kompilieren. Wenn Sie die in Python verfügbaren Operatoren definieren möchten, können Sie Ihre Registrierungen auf ähnliche Weise in den Python-Wrapper-Code einfügen.

Zur Unterstützung einer Reihe von Vorgängen kann ein ähnlicher Prozess wie oben verwendet werden, anstatt nur einen einzelnen Operator. Sie können einfach so viele AddCustom-Operatoren hinzufügen, wie Sie benötigen. Außerdem können Sie mit MutableOpResolver Implementierungen von builtins mithilfe von AddBuiltin überschreiben.

Anbieter testen und Profil erstellen

Wenn Sie ein Profil des Vorgangs mit dem Benchmark-Tool von TensorFlow Lite erstellen möchten, können Sie das Benchmark-Modelltool für TensorFlow Lite verwenden. Zu Testzwecken können Sie Ihren lokalen Build von TensorFlow Lite auf Ihre benutzerdefinierte Operation aufmerksam machen, indem Sie register.cc den entsprechenden AddCustom-Aufruf (wie oben gezeigt) hinzufügen.

Best Practices

Gehen Sie beim Optimieren von Arbeitsspeicherzuweisungen und -entfernungen vorsichtig vor. Die Zuweisung von Arbeitsspeicher in Prepare ist effizienter als in Invoke und die Zuweisung von Arbeitsspeicher vor einer Schleife ist besser als in jedem Durchlauf. Verwenden Sie temporäre Tensordaten, anstatt sich selbst zu definieren (siehe Punkt 2). Verwenden Sie Zeiger/Verweise, anstatt so viel wie möglich zu kopieren.

Wenn eine Datenstruktur während des gesamten Vorgangs bestehen bleibt, empfehlen wir, den Speicher mithilfe temporärer Tensoren vorab zuzuordnen. Möglicherweise müssen Sie eine OpData-Struktur verwenden, um in anderen Funktionen auf die Tensorindizes zu verweisen. Sehen Sie sich dazu das Beispiel im Kernel für Faltung an. Unten sehen Sie ein Beispiel für ein Code-Snippet.

struct MyOpData {
  int temp_tensor_index;
  ...
};

void* Init(TfLiteOpaqueContext* context,
    const char* buffer, size_t length) {
  auto* op_data = new MyOpData{};
  ...
  return op_data;
}
void Free(TfLiteOpaqueContext* context, void* buffer) {
  ...
  delete reinterpret_cast<MyOpData*>(buffer);
}
TfLiteStatus Prepare(TfLiteOpaqueContext* context,
                     TfLiteOpaqueNode* node) {
  ...
  auto* op_data =
      reinterpret_cast<MyOpData*>(TfLiteOpaqueNodeGetUserData(node));
  const int num_temporaries = 1;
  int temporary_tensor_indices[num_temporaries];
  TfLiteOpaqueTensorBuilder* builder = TfLiteOpaqueTensorBuilderCreate();
  TfLiteOpaqueTensorBuilderSetType(builder, kTfLiteFloat32);
  TfLiteOpaqueTensorBuilderSetAllocationType(builder, kTfLiteArenaRw);
  TfLiteOpaqueContextAddTensor(context, builder,
      &temporary_tensor_indices[0]);
  TfLiteOpaqueTensorBuilderDelete(builder);
  TfLiteOpaqueNodeSetTemporaries(node, temporary_tensor_indices,
      num_temporaries);
  op_data->temp_tensor_index = temporary_tensor_indices[0];
  ...
  return kTfLiteOk;
}
TfLiteStatus Invoke(TfLiteOpaqueContext* context,
                    TfLiteOpaqueNode* node) {
  ...
  auto* op_data = reinterpret_cast<MyOpData*>(
      TfLiteOpaqueNodeGetUserData(node));
  TfLiteOpaqueTensor* temp_tensor =
      TfLiteOpaqueContextGetOpaqueTensor(context,
          op_data->temp_tensor_index);
  TF_LITE_OPAQUE_ENSURE(context,
      TfLiteTensorType(temp_tensor) == kTfLiteFloat32);
  TF_LITE_OPAQUE_ENSURE(context,
      TfLiteTensorGetAllocationType(temp_Tensor) == kTfLiteArenaRw);
  void *temp_data = TfLiteTensorData(temp_tensor);
  TF_LITE_OPAQUE_ENSURE(context, temp_data != nullptr);
  ...
  return kTfLiteOk;
}

Wenn Sie damit nicht zu viel Arbeitsspeicher verschwenden, sollten Sie ein statisches Array mit fester Größe (oder eine vorab zugewiesene std::vector in Resize) verwenden, anstatt bei jeder Ausführungsdurchführung eine dynamisch zugewiesene std::vector zu verwenden.
Instanziieren Sie keine Containervorlagen für Standardbibliotheken, die noch nicht vorhanden sind, da sie sich auf die Binärgröße auswirken. Wenn Sie in Ihrem Vorgang beispielsweise ein std::map benötigen, das in anderen Kerneln nicht vorhanden ist, könnte die Verwendung eines std::vector mit direkter Indexierungszuordnung funktionieren, während die Binärgröße klein bleibt. Sehen Sie sich an, was andere Kernel verwenden, um Erkenntnisse zu gewinnen (oder zu fragen).
Prüfen Sie den Zeiger auf den von malloc zurückgegebenen Arbeitsspeicher. Wenn dieser Zeiger nullptr ist, sollten keine Vorgänge mit diesem Zeiger ausgeführt werden. Wenn in einer Funktion malloc angegeben wird und ein Fehler beendet wird, heben Sie die Zuweisung des Arbeitsspeichers auf, bevor Sie die Funktion beenden.
Verwenden Sie TF_LITE_OPAQUE_ENSURE(context, condition), um nach einer bestimmten Bedingung zu suchen. Ihr Code darf den Arbeitsspeicher nicht hängen lassen, wenn TF_LITE_OPAQUE_ENSURE verwendet wird. Das heißt, diese Makros sollten verwendet werden, bevor Ressourcen zugewiesen werden, die Speicherlecks verursachen.