Operadores personalizados

Como a biblioteca de operadores integrada do TensorFlow Lite oferece suporte apenas a um número limitado de operadores do TensorFlow, nem todos os modelos são conversíveis. Para obter detalhes, consulte a compatibilidade do operador .

Para permitir a conversão, os usuários podem fornecer sua própria implementação personalizada de um operador TensorFlow sem suporte no TensorFlow Lite, conhecido como operador personalizado. Se, em vez disso, você quiser combinar uma série de operadores TensorFlow sem suporte (ou com suporte) em um único operador personalizado otimizado e fundido, consulte Fusão do operador .

O uso de operadores personalizados consiste em quatro etapas.

Vamos percorrer um exemplo completo de execução de um modelo com um operador personalizado tf.atan (denominado Atan , consulte #create_a_tensorflow_model) que é compatível com o TensorFlow, mas não é compatível com o TensorFlow Lite.

O operador TensorFlow Text é um exemplo de operador personalizado. Veja o tutorial Convert TF Text to TF Lite para um exemplo de código.

Exemplo: operador Atan personalizado

Vejamos um exemplo de suporte a um operador do TensorFlow que o TensorFlow Lite não possui. Suponha que estamos usando o operador Atan e que estamos construindo um modelo muito simples para uma função y = atan(x + offset) , onde offset pode ser treinado.

Criar um modelo do TensorFlow

O trecho de código a seguir treina um modelo simples do TensorFlow. Este modelo contém apenas um operador personalizado chamado Atan , que é uma função y = atan(x + offset) , onde 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 do TensorFlow Lite com os sinalizadores do conversor padrão, receberá a seguinte mensagem de erro:

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

Converter para um modelo do TensorFlow Lite

Crie um modelo do TensorFlow Lite 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ê executá-lo com o interpretador padrão usando comandos como a seguir:

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

Você ainda receberá o erro:

Encountered unresolved custom op: Atan.

Crie e registre o operador.

Todos os operadores do TensorFlow Lite (personalizados e integrados) são definidos usando uma interface simples de C puro que consiste em quatro funções:

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;

Consulte common.h para obter detalhes sobre TfLiteContext e TfLiteNode . O primeiro fornece recursos de relatório de erros e acesso a objetos globais, incluindo todos os tensores. O último permite que as implementações acessem suas entradas e saídas.

Quando o interpretador carrega um modelo, ele chama init() uma vez para cada nó no grafo. 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 os nomes dos parâmetros para seus valores. O buffer está vazio para operações integradas porque o interpretador já analisou os parâmetros de operação. As implementações de kernel que requerem estado devem inicializá-lo aqui e transferir a propriedade para o chamador. Para cada chamada init() , haverá uma chamada correspondente para free() , permitindo que as implementações descartem o buffer que podem ter alocado em init() .

Sempre que os tensores de entrada forem redimensionados, o interpretador percorrerá o gráfico notificando as implementações da alteração. Isso lhes dá a chance de redimensionar seu buffer interno, verificar a validade das formas e tipos de entrada e recalcular as formas de saída. Tudo isso é feito por meio de prepare() , e as implementações podem acessar seu estado usando node->user_data .

Finalmente, cada vez que a inferência é executada, o interpretador percorre o grafo chamando invoke() , e aqui também o estado está disponível como node->user_data .

As operações personalizadas podem ser implementadas exatamente da mesma maneira que as operações integradas, definindo essas quatro funções e uma função de registro global que geralmente se parece com isso:

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

Observe que o registro não é automático e uma chamada explícita para Register_MY_CUSTOM_OP deve ser feita. Enquanto o BuiltinOpResolver padrão (disponível no destino :builtin_ops ) cuida do registro de builtins, as operações personalizadas terão que ser coletadas em bibliotecas personalizadas separadas.

Definindo o kernel no tempo de execução do TensorFlow Lite

Tudo o que precisamos fazer para usar o op no TensorFlow Lite é definir duas funções ( Prepare e Eval ) e construir um TfLiteRegistration :

TfLiteStatus AtanPrepare(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 AtanEval(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  const TfLiteTensor* input = GetInput(context, node, 0);
  TfLiteTensor* output = GetOutput(context, node, 0);

  float* input_data = GetTensorData<float>(input);
  float* output_data = GetTensorData<float>(output);

  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] = atan(input_data[i]);
  }
  return kTfLiteOk;
}

TfLiteRegistration* Register_ATAN() {
  static TfLiteRegistration r = {nullptr, nullptr, AtanPrepare, AtanEval};
  return &r;
}

Ao inicializar o OpResolver , adicione o op personalizado ao resolvedor (veja abaixo um exemplo). Isso registrará o operador no Tensorflow Lite para que o TensorFlow Lite possa usar a nova implementação. Observe que os dois últimos argumentos em TfLiteRegistration correspondem às funções AtanPrepare e AtanEval que você definiu para a operação personalizada. Se você usou as funções AtanInit e AtanFree para inicializar variáveis ​​usadas no op e para liberar espaço, respectivamente, elas seriam adicionadas aos dois primeiros argumentos de TfLiteRegistration ; esses argumentos são definidos como nullptr neste exemplo.

Registre o operador com a biblioteca do kernel

Agora precisamos registrar o operador com a biblioteca do kernel. Isso é feito com um OpResolver . Nos bastidores, o interpretador carregará uma biblioteca de kernels que serão designados para executar cada um dos operadores do modelo. Embora a biblioteca padrão contenha apenas kernels integrados, é possível substituí-la/aumentá-la por uma biblioteca personalizada de operadores operacionais.

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

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

O uso regular requer que você use o BuiltinOpResolver e escreva:

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

Para adicionar a operação personalizada criada acima, chame AddOp (antes de passar o resolvedor para o InterpreterBuilder ):

resolver.AddCustom("Atan", Register_ATAN());

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

Se quiser definir seus operadores customizados em Java, você precisaria construir sua própria camada JNI customizada e compilar seu próprio AAR neste código jni . Da mesma forma, se você deseja definir esses operadores disponíveis no Python, pode colocar seus registros no código wrapper do Python .

Observe que um processo semelhante ao anterior pode ser seguido para suportar um conjunto de operações em vez de um único operador. Basta adicionar quantos operadores AddCustom forem necessários. Além disso, BuiltinOpResolver também permite substituir implementações de builtins usando o AddBuiltin .

Teste e crie o perfil do seu operador

Para criar o perfil de sua operação com a ferramenta de benchmark TensorFlow Lite, você pode usar a ferramenta de modelo de benchmark para TensorFlow Lite. Para fins de teste, você pode tornar sua compilação local do TensorFlow Lite ciente de sua operação personalizada adicionando a chamada AddCustom apropriada (conforme mostrado acima) a register.cc

Melhores Práticas

  1. Otimize as alocações e desalocações de memória com cautela. 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 malocar você mesmo (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, aconselhamos pré-alocar a memória usando tensores temporários. Você pode precisar usar a estrutura OpData para referenciar os índices do tensor em outras funções. Veja o exemplo no kernel para convolução . Um trecho de código de amostra está abaixo

    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. 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 a cada iteração de execução.

  4. Evite instanciar modelos de contêiner de biblioteca padrão que ainda não existam, porque eles afetam o tamanho binário. Por exemplo, se você precisar de um std::map em sua operação que não existe em outros kernels, usar um std::vector com mapeamento de indexação direta pode funcionar enquanto mantém o tamanho do binário pequeno. Veja o que outros kernels usam para obter informações (ou pergunte).

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

  6. Use TF_LITE_ENSURE(context, condition) para verificar uma condição específica. Seu código não deve deixar memória travada quando TF_LITE_ENSURE for usado, ou seja, essas macros devem ser usadas antes que qualquer recurso seja alocado que irá vazar.