Operatori personalizzati

Poiché la libreria di operatori integrati LiteRT supporta solo un numero limitato di operatori TensorFlow, non tutti i modelli sono convertibili. Per maggiori dettagli, consulta la compatibilità con gli operatori.

Per consentire la conversione, gli utenti possono fornire la propria implementazione personalizzata di un operatore TensorFlow non supportato in LiteRT, noto come operatore personalizzato. Se invece vuoi combinare una serie di operatori TensorFlow non supportati (o supportati) in un unico operatore personalizzato ottimizzato e unito, consulta la sezione Unione di operatori.

L'utilizzo di operatori personalizzati prevede quattro passaggi.

Vediamo un esempio end-to-end di esecuzione di un modello con un operatore personalizzato tf.atan (denominato Atan, consulta Creare un modello TensorFlow) supportato in TensorFlow, ma non in LiteRT.

L'operatore TensorFlow Text è un esempio di operatore personalizzato. Per un esempio di codice, consulta il tutorial Convertire TF Text in LiteRT.

Esempio: operatore Atan personalizzato

Vediamo un esempio di supporto di un operatore TensorFlow che LiteRT non ha. Supponiamo di utilizzare l'operatore Atan e di creare un modello molto semplice per una funzione y = atan(x + offset), dove offset è addestrabile.

Crea un modello TensorFlow

Il seguente snippet di codice addestra un semplice modello TensorFlow. Questo modello contiene solo un operatore personalizzato denominato Atan, che è una funzione y = atan(x + offset), dove offset è addestrabile.

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

A questo punto, se provi a generare un modello LiteRT con i flag del convertitore predefiniti, riceverai il seguente messaggio di errore:

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

Convertire in un modello LiteRT

Crea un modello LiteRT con operatori personalizzati impostando l'attributo del convertitore allow_custom_ops come mostrato di seguito:

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

A questo punto, se lo esegui con l'interprete predefinito utilizzando comandi come segue:

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

Continuerai a ricevere l'errore:

Encountered unresolved custom op: Atan.

Crea e registra l'operatore.

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

Gli operatori personalizzati LiteRT sono definiti utilizzando una semplice API pure-C che consiste in un tipo opaco (TfLiteOperator) e funzioni correlate.

TfLiteOperator è un tipo opaco:

typedef struct TfLiteOperator TfLiteOperator;

TfLiteOperator memorizza l'identità e l'implementazione dell'operatore. Tieni presente che l'operatore è distinto dai relativi operandi, che sono memorizzati nei nodi del grafico LiteRT per i nodi che chiamano l'operatore.

Le istanze di questo tipo vengono create con chiamate a TfLiteOperatorCreate e possono essere distrutte chiamando TfLiteOperatorDelete.

L'identità dell'operatore viene impostata tramite i parametri della funzione costruttore TfLiteOperatorCreate:

TfLiteOperator*
TfLiteOperatorCreate(
    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.
);

L'implementazione dell'operatore può definire "metodi" con le seguenti firme. Tutti questi metodi sono facoltativi, ma affinché un operatore venga valutato correttamente, l'implementazione dell'operatore deve definire e impostare (utilizzando le funzioni setter) almeno i metodi Prepare e Invoke.

// 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);

I nomi delle funzioni (o i prefissi dello spazio dei nomi, per C++) nell'implementazione dell'operazione non devono corrispondere ai nomi delle funzioni nel precedente snippet di codice, poiché l'API TF Lite custom ops utilizzerà solo i relativi indirizzi. Infatti, ti consigliamo di dichiararli in uno spazio dei nomi anonimo o come funzioni statiche.

Tuttavia, è consigliabile includere il nome dell'operatore come spazio dei nomi o prefisso in questi nomi di funzioni:

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.
      

Poiché si tratta di un'API C, questi "metodi" vengono implementati come puntatori a funzioni C nel tipo TfLiteOperator, che vengono impostati passando gli indirizzi delle funzioni di implementazione alle funzioni setter corrispondenti TfLiteOperatorSetMethodName:

void TfLiteOperatorSetInit(
    TfLiteOperator* operator,
    void* (*init)(TfLiteOpaqueContext* context, const char* buffer,
                  size_t length));
void TfLiteOperatorSetFree(
    TfLiteOperator* operator,
    void (*free)(TfLiteOpaqueContext* context, void* data));
void TfLiteOperatorSetPrepare(
    TfLiteOperator* operator,
    TfLiteStatus (*prepare)(TfLiteOpaqueContext* context,
                            TfLiteOpaqueNode* node));
void TfLiteOperatorSetInvoke(
    TfLiteOperator* operator,
    TfLiteStatus (*invoke)(TfLiteOpaqueContext* context,
                           TfLiteOpaqueNode* node));
void TfLiteOperatorSetAsyncKernel(
    TfLiteOperator* operator,
    struct TfLiteAsyncKernel* (*async_kernel)(TfLiteOpaqueContext* context,
                                              TfLiteOpaqueNode* node));

Per informazioni dettagliate su TfLiteContext e TfLiteNode, consulta common.h. TfLiteContext fornisce funzionalità di segnalazione degli errori e accesso agli oggetti globali, inclusi tutti i tensori. TfLiteNode consente alle implementazioni degli operatori di accedere ai propri input e output.

Quando l'interprete carica un modello, chiama il metodo Init() una volta per ogni nodo del grafico. Un determinato Init() verrà chiamato più di una volta se l'operazione viene utilizzata più volte nel grafico. Per le operazioni personalizzate verrà fornito un buffer di configurazione contenente un flexbuffer che mappa i nomi dei parametri ai relativi valori. Il buffer è vuoto per le operazioni integrate perché l'interprete ha già analizzato i parametri dell'operazione. Le implementazioni del kernel che richiedono lo stato devono inizializzarlo qui e trasferire la proprietà al chiamante. Per ogni chiamata Init(), ci sarà una chiamata corrispondente a Free(), che consentirà alle implementazioni di eliminare il buffer che potrebbero aver allocato in Init().

Ogni volta che le dimensioni dei tensori di input vengono modificate, l'interprete esamina il grafico notificando le implementazioni della modifica. In questo modo hanno la possibilità di ridimensionare il buffer interno, controllare la validità di forme e tipi di input e ricalcolare le forme di output. Tutto questo viene fatto tramite il metodo Prepare() e le implementazioni possono accedere al proprio stato utilizzando TfLiteOpaqueNodeGetUserData(node).

Infine, ogni volta che viene eseguita l'inferenza, l'interprete attraversa il grafico chiamando il metodo Invoke() e anche qui lo stato è disponibile come TfLiteOpaqueNodeGetUserData(node).

Le operazioni personalizzate possono essere implementate definendo le funzioni "method" e poi definendo una funzione che restituisce un'istanza di TfLiteOperator costruita chiamando TfLiteOperatorCreate e poi i metodi setter pertinenti:

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 TfLiteOperator* MyCustomOperator() {
    // Singleton instance, intentionally never destroyed.
    static const TfLiteOperator* my_custom_op = ()[] {
        TfLiteOperator* r =
            TfLiteOperatorCreate(
                kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1);
        TfLiteOperatorSetInit(r, Init);
        TfLiteOperatorSetFree(r, Free);
        TfLiteOperatorSetPrepare(r, Prepare);
        TfLiteOperatorSetInvoke(r, Eval);
        return r;
      };
    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 TfLiteOperator* MyCustomOpCreate() {
  const TfLiteOperator* r =
      TfLiteOperatorCreate(
          kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1);
  TfLiteOperatorSetInit(r, MyCustomOpInit);
  TfLiteOperatorSetFree(r, MyCustomOpFree);
  TfLiteOperatorSetPrepare(r, MyCustomOpPrepare);
  TfLiteOperatorSetInvoke(r, MyCustomOpEval);
  return r;
}

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

Tieni presente che la registrazione non è automatica e deve essere effettuata una chiamata esplicita alla tua funzione MyCustomOperator (vedi i dettagli di seguito). Mentre l'BuiltinOpResolver standard (disponibile dalla destinazione :builtin_ops) si occupa della registrazione dei componenti integrati, le operazioni personalizzate dovranno essere raccolte in librerie personalizzate separate.

Definizione del kernel nel runtime LiteRT

Per utilizzare l'operatore in LiteRT, è sufficiente definire due funzioni (Prepare e Eval) e una terza per costruire un TfLiteOperator:

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<float*>(TfLiteOpaqueTensorData(input));
      float* output_data = static_cast<float*>(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 TfLiteOperator* AtanOperator() {
    // Singleton instance, intentionally never destroyed.
    static const TfLiteOperator* atan_op = ()[] {
        auto* r = TfLiteOperatorCreate(
            kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1);
        TfLiteOperatorSetPrepare(r, Prepare);
        TfLiteOperatorSetInvoke(r, Eval);
        return r;
      };
    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<float*>(TfLiteOpaqueTensorData(input));
  float* output_data = static_cast<float*>(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 TfLiteOperator* AtanOpCreate() {
  TfLiteOperator* r = TfLiteOperatorCreate(
          kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1);
  TfLiteOperatorSetPrepare(r, Prepare);
  TfLiteOperatorSetInvoke(r, Eval);
  return r;
}

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

Quando inizializzi OpResolver, aggiungi l'operazione personalizzata al resolver (vedi l'esempio di seguito). In questo modo, l'operatore verrà registrato con LiteRT in modo che LiteRT possa utilizzare la nuova implementazione.

Registra l'operatore con la libreria del kernel

Ora dobbiamo registrare l'operatore con la libreria del kernel. Questa operazione viene eseguita con un OpResolver. Dietro le quinte, l'interprete caricherà una libreria di kernel che verranno assegnati per eseguire ciascuno degli operatori nel modello. Sebbene la libreria predefinita contenga solo kernel integrati, è possibile sostituirla/integrarla con una libreria personalizzata di operatori op.

La classe OpResolver, che traduce i codici e i nomi degli operatori in codice effettivo, è definita come segue:

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

Tieni presente che, per la compatibilità con le versioni precedenti, questa classe utilizza il tipo concreto precedente TfLiteRegistration anziché il tipo opaco TfLiteOperator, ma lo struct TfLiteRegistration contiene un campo registration_external di tipo TfLiteOperator*.

Le classi MutableOpResolver e BuiltinOpResolver derivano da OpResolver:

class MutableOpResolver : public OpResolver {
 public:
  MutableOpResolver();  // Constructs an initially empty op resolver.
  void AddAll(const MutableOpResolver& other);
  ...
};

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

L'utilizzo regolare (senza operazioni personalizzate) richiede l'utilizzo di BuiltinOpResolver e la scrittura di:

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

Per aggiungere l'operazione personalizzata creata sopra, puoi utilizzare invece un MutableOpResolver, e chiamare tflite::AddOp (prima di passare il resolver a InterpreterBuilder):

tflite::ops::builtin::MutableOpResolver resolver;
resolver.AddAll(tflite::ops::builtin::BuiltinOpResolver());
tflite::AddOp(&resolver, AtanOpRegistration());

Se il set di operazioni integrate viene considerato troppo grande, è possibile generare automaticamente un nuovo OpResolver in base a un determinato sottoinsieme di operazioni, possibilmente solo quelle contenute in un determinato modello. Si tratta dell'equivalente della registrazione selettiva di TensorFlow (e una versione semplice è disponibile nella directory tools).

Se vuoi definire i tuoi operatori personalizzati in Java, al momento devi creare il tuo livello JNI personalizzato e compilare il tuo AAR in questo codice JNI. Allo stesso modo, se vuoi definire questi operatori disponibili in Python, puoi inserire le registrazioni nel codice wrapper Python.

Tieni presente che è possibile seguire una procedura simile a quella descritta sopra per supportare un insieme di operazioni anziché un singolo operatore. Aggiungi tutti gli operatori AddCustom che ti servono. Inoltre, MutableOpResolver ti consente anche di eseguire l'override delle implementazioni dei componenti integrati utilizzando AddBuiltin.

Testare e profilare l'operatore

Per profilare l'operazione con lo strumento di benchmark LiteRT, puoi utilizzare lo strumento per il modello di benchmark per LiteRT. A scopo di test, puoi rendere la tua build locale di LiteRT consapevole della tua operazione personalizzata aggiungendo la chiamata AddCustom appropriata (come mostrato sopra) a register.cc

Best practice

  1. Ottimizza le allocazioni e le deallocazioni di memoria con cautela. L'allocazione della memoria in Prepare è più efficiente rispetto a Invoke e l'allocazione della memoria prima di un ciclo è migliore rispetto a ogni iterazione. Utilizza i dati dei tensori temporanei anziché allocarli manualmente (vedi punto 2). Utilizza puntatori/riferimenti anziché copiare il più possibile.

  2. Se una struttura di dati persiste durante l'intera operazione, ti consigliamo di preallocare la memoria utilizzando tensori temporanei. Potresti dover utilizzare una struct OpData per fare riferimento agli indici dei tensori in altre funzioni. Consulta l'esempio nel kernel per la convoluzione. Di seguito è riportato un esempio di snippet di codice.

    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;
    }
    
  3. Se non comporta un eccessivo spreco di memoria, preferisci utilizzare un array statico di dimensioni fisse (o un std::vector preallocato in Resize) anziché utilizzare un std::vector allocato dinamicamente a ogni iterazione di esecuzione.

  4. Evita di creare istanze di modelli di contenitori della libreria standard che non esistono già, perché influiscono sulle dimensioni del file binario. Ad esempio, se hai bisogno di un std::map nella tua operazione che non esiste in altri kernel, l'utilizzo di un std::vector con mappatura dell'indicizzazione diretta potrebbe funzionare mantenendo le dimensioni del file binario ridotte. Scopri cosa usano gli altri kernel per ottenere informazioni (o chiedi).

  5. Controlla il puntatore alla memoria restituita da malloc. Se questo puntatore è nullptr, non devono essere eseguite operazioni che lo utilizzano. Se malloc in una funzione e hai un'uscita di errore, dealloca la memoria prima di uscire.

  6. Utilizza TF_LITE_OPAQUE_ENSURE(context, condition) per verificare una condizione specifica. Il codice non deve lasciare la memoria in sospeso quando viene utilizzato TF_LITE_OPAQUE_ENSURE, ovvero queste macro devono essere utilizzate prima che vengano allocate risorse che verranno perse.