Operadores personalizados

Como a biblioteca de operadores integrados do LiteRT só é compatível com um número limitado de operadores do TensorFlow, nem todos os modelos podem ser convertidos. Para mais detalhes, consulte compatibilidade de operadores.

Para permitir a conversão, os usuários podem fornecer a própria implementação personalizada de um operador do TensorFlow não compatível no LiteRT, conhecido como operador personalizado. Se você quiser combinar uma série de operadores do TensorFlow não compatíveis (ou compatíveis) em um único operador personalizado otimizado e fundido, consulte fusão de operadores.

O uso de operadores personalizados consiste em quatro etapas.

Vamos analisar um exemplo de ponta a ponta de execução de um modelo com um operador personalizado tf.atan (chamado de Atan.Consulte Criar um modelo do TensorFlow), que é compatível com o TensorFlow, mas não com o LiteRT.

O operador do TensorFlow Text é um exemplo de operador personalizado. Consulte o tutorial Converter TF Text para LiteRT (em inglês) para ver um exemplo de código.

Exemplo: operador Atan personalizado

Vamos analisar um exemplo de suporte a um operador do TensorFlow que a LiteRT não tem. Suponha que estamos usando o operador Atan e criando um modelo muito simples para uma função y = atan(x + offset), em que offset pode ser treinado.

Criar um modelo do TensorFlow

O snippet de código a seguir treina um modelo simples do TensorFlow. Esse modelo contém apenas um operador personalizado chamado Atan, que é uma função y = atan(x + offset), em que offset pode ser treinado.

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

Neste ponto, se você tentar gerar um modelo LiteRT com as flags de conversor padrão, vai receber a seguinte mensagem de erro:

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

Converter para um modelo LiteRT

Crie um modelo LiteRT com operadores personalizados definindo o atributo do conversor allow_custom_ops, conforme mostrado abaixo:

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

Neste ponto, se você executar com o interpretador padrão usando comandos como a seguir:

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

Você ainda vai receber o erro:

Encountered unresolved custom op: Atan.

Crie e registre o operador.

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

Os operadores personalizados do LiteRT são definidos usando uma API simples de C puro que consiste em um tipo opaco (TfLiteOperator) e funções relacionadas.

TfLiteOperator é um tipo opaco:

typedef struct TfLiteOperator TfLiteOperator;

TfLiteOperator armazena a identidade e a implementação do operador. O operador é diferente dos operandos, que são armazenados nos nós do gráfico LiteRT para nós que chamam o operador.

As instâncias desse tipo são construídas com chamadas para TfLiteOperatorCreate e podem ser destruídas chamando TfLiteOperatorDelete.

A identidade do operador é definida pelos parâmetros da função construtora 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.
);

A implementação do operador pode definir "métodos" com as seguintes assinaturas. Todos esses métodos são opcionais, mas, para que um operador seja avaliado com sucesso, a implementação dele precisa definir e definir (usando as funções setter) pelo menos os métodos 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);

A função names (ou prefixos de namespace, para C++) na sua implementação de operação não precisa corresponder aos nomes de função no snippet de código acima, já que a API de operações personalizadas do TF Lite só usará os endereços delas. Recomendamos que você declare essas funções em um namespace anônimo ou como funções estáticas.

No entanto, é recomendável incluir o nome do operador como um namespace ou prefixo nesses nomes de função:

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.

Como essa é uma API em C, esses "métodos" são implementados como ponteiros de função em C no tipo TfLiteOperator, que são definidos transmitindo os endereços das funções de implementação para as funções setter correspondentes 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));

Consulte common.h para mais detalhes sobre TfLiteContext e TfLiteNode. TfLiteContext oferece recursos de relatórios de erros e acesso a objetos globais, incluindo todos os tensores. O TfLiteNode permite que as implementações de operadores acessem as entradas e saídas.

Quando o intérprete carrega um modelo, ele chama o método Init() uma vez para cada nó no gráfico. Um determinado Init() será chamado mais de uma vez se a operação for usada várias vezes no gráfico. Para operações personalizadas, um buffer de configuração será fornecido, contendo um flexbuffer que mapeia nomes de parâmetros para os valores deles. O buffer está vazio para operações integradas porque o interpretador já analisou os parâmetros da operação. As implementações do kernel que exigem estado precisam inicializá-lo aqui e transferir a propriedade para o autor da chamada. Para cada chamada de Init(), haverá uma chamada correspondente para Free(), permitindo que as implementações descartem o buffer que pode ter sido alocado em Init().

Sempre que os tensores de entrada forem redimensionados, o intérprete vai passar pelo gráfico notificando as implementações da mudança. Isso dá a eles a chance de redimensionar o buffer interno, verificar a validade de formas e tipos de entrada e recalcular as formas de saída. Tudo isso é feito com o método Prepare(), e as implementações podem acessar o estado usando TfLiteOpaqueNodeGetUserData(node).

Por fim, cada vez que a inferência é executada, o intérprete percorre o gráfico chamando o método Invoke(). Aqui também o estado está disponível como TfLiteOpaqueNodeGetUserData(node).

As operações personalizadas podem ser implementadas definindo essas funções de "método" e, em seguida, definindo uma função que retorna uma instância de TfLiteOperator construída chamando TfLiteOperatorCreate e, em seguida, os métodos setter relevantes:

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;
}

O registro não é automático. É necessário fazer uma chamada explícita para a função MyCustomOperator (confira os detalhes abaixo). Embora o BuiltinOpResolver padrão (disponível na meta :builtin_ops) cuide do registro de builtins, as operações personalizadas precisam ser coletadas em bibliotecas personalizadas separadas.

Como definir o kernel no ambiente de execução do LiteRT

Para usar a operação no LiteRT, basta definir duas funções (Prepare e Eval) e uma terceira para construir um 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;
}

Ao inicializar o OpResolver, adicione a operação personalizada ao resolver. Consulte um exemplo abaixo. Isso vai registrar o operador com o LiteRT para que ele possa usar a nova implementação.

Registre o operador com a biblioteca do kernel

Agora precisamos registrar o operador na biblioteca do kernel. Isso é feito com um OpResolver. Nos bastidores, o intérprete carrega uma biblioteca de kernels que são atribuídos para executar cada um dos operadores no modelo. Embora a biblioteca padrão contenha apenas kernels integrados, é possível substituí-la/aumentá-la com operadores de biblioteca personalizados.

A classe OpResolver, que traduz códigos e nomes de operadores em código real, é definida assim:

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

Para oferecer compatibilidade com versões anteriores, essa classe usa o tipo concreto mais antigo TfLiteRegistration em vez do tipo opaco TfLiteOperator, mas a struct TfLiteRegistration contém um campo registration_external do tipo TfLiteOperator*.

As classes MutableOpResolver e BuiltinOpResolver são derivadas de 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.
};

O uso normal (sem operações personalizadas) exige que você use BuiltinOpResolver e escreva:

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

Para adicionar a operação personalizada criada acima, use um MutableOpResolver e chame tflite::AddOp antes de transmitir o resolver para o InterpreterBuilder:

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

Se o conjunto de operações integradas for considerado muito grande, um novo OpResolver poderá ser gerado por código com base em um determinado subconjunto de operações, possivelmente apenas as contidas em um determinado modelo. Isso é o equivalente ao registro seletivo do TensorFlow. Uma versão simples dele está disponível no diretório tools.

Se você quiser definir seus operadores personalizados em Java, precisará criar sua própria camada JNI personalizada e compilar seu próprio AAR neste código JNI. Da mesma forma, se você quiser definir esses operadores disponíveis em Python, coloque os registros no código do wrapper Python.

Um processo semelhante ao acima pode ser seguido para oferecer suporte a um conjunto de operações em vez de um único operador. Basta adicionar quantos operadores AddCustom forem necessários. Além disso, MutableOpResolver também permite substituir implementações de builtins usando AddBuiltin.

Testar e criar perfil do operador

Para criar um perfil da sua operação com a ferramenta de comparativo de mercado do LiteRT, use a ferramenta de modelo de comparativo de mercado para o LiteRT. Para fins de teste, adicione a chamada AddCustom apropriada (como mostrado acima) a register.cc para que seu build local do LiteRT reconheça a operação personalizada.

Práticas recomendadas

  1. Otimize as alocações e desalocações de memória com cuidado. Alocar memória em Prepare é mais eficiente do que em Invoke, e alocar memória antes de um loop é melhor do que em cada iteração. Use dados de tensores temporários em vez de fazer malloc (consulte o item 2). Use ponteiros/referências em vez de copiar o máximo possível.

  2. Se uma estrutura de dados persistir durante toda a operação, recomendamos pré-alocar a memória usando tensores temporários. Talvez seja necessário usar uma struct OpData para referenciar os índices de tensor em outras funções. Confira o exemplo no kernel para convolução. Confira abaixo um exemplo de snippet de código.

    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 não custar muita memória desperdiçada, prefira usar uma matriz estática de tamanho fixo (ou um std::vector pré-alocado em Resize) em vez de usar um std::vector alocado dinamicamente em cada iteração de execução.

  4. Evite instanciar modelos de contêiner da biblioteca padrão que ainda não existem, porque eles afetam o tamanho do binário. Por exemplo, se você precisar de um std::map na sua operação que não exista em outros kernels, usar um std::vector com mapeamento de indexação direta pode funcionar e manter o tamanho binário pequeno. Veja o que outros kernels usam para ter insights (ou pergunte).

  5. Verifique o ponteiro para a memória retornada por malloc. Se esse ponteiro for nullptr, nenhuma operação poderá ser realizada usando esse ponteiro. Se você malloc em uma função e tiver uma saída de erro, desalocar a memória antes de sair.

  6. Use TF_LITE_OPAQUE_ENSURE(context, condition) para verificar uma condição específica. Seu código não pode deixar a memória pendente quando TF_LITE_OPAQUE_ENSURE é usado. Ou seja, essas macros precisam ser usadas antes que qualquer recurso seja alocado e cause um vazamento.