Operadores personalizados

Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.

Como a biblioteca de operadores integrada do TensorFlow Lite é compatível apenas com um número limitado de operadores do TensorFlow, nem todos os modelos são conversíveis. Para obter detalhes, consulte 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ê desejar combinar uma série de operadores TensorFlow não suportados (ou suportados) em um único operador personalizado otimizado e fundido, consulte a fusão de operadores .

O uso de operadores personalizados consiste em quatro etapas.

Vejamos um exemplo completo de execução de um modelo com um operador personalizado tf.sin (chamado Sin , 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. Consulte o tutorial Converter texto do TF em TF Lite para obter um exemplo de código.

Exemplo: operador Sin personalizado

Vejamos um exemplo de suporte a um operador do TensorFlow que o TensorFlow Lite não possui. Suponha que estamos usando o operador Sin e que estamos construindo um modelo muito simples para uma função y = sin(x + offset) , onde offset é treinável.

Criar um modelo do TensorFlow

O snippet de código a seguir treina um modelo simples do TensorFlow. Este modelo contém apenas um operador personalizado chamado Sin , que é uma função y = sin(x + offset) , onde o offset é treinável.

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

Nesse ponto, se você tentar gerar um modelo do TensorFlow Lite com os sinalizadores padrão do conversor, receberá a seguinte mensagem de erro:

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.

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([sin.get_concrete_function(x)], sin)
converter.allow_custom_ops = True
tflite_model = converter.convert()

Neste ponto, se você executá-lo com o interpretador padrão, você receberá as seguintes mensagens de erro:

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

Crie e registre o operador.

Todos os operadores do TensorFlow Lite (personalizados e integrados) são definidos usando uma interface C puro simples 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. Este ú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 gráfico. Um dado init() será chamado mais de uma vez se o op for usado 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 internas porque o interpretador já analisou os parâmetros de operação. As implementações de kernel que exigem 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 passará pelo gráfico notificando as implementações da mudança. 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. Isso tudo é feito através 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 gráfico 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 internas, 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 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;
}

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 SinPrepare e SinEval que você definiu para a operação personalizada. Se você usou as funções SinInit e SinFree para inicializar as 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 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 operacionais de biblioteca personalizados.

A classe OpResolver , que traduz códigos e nomes de operadores 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 o op personalizado criado acima, você chama AddOp (antes de passar o resolvedor para o InterpreterBuilder ):

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

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 de tools ).

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

Observe que um processo semelhante ao acima pode ser seguido para suportar um conjunto de operações em vez de um único operador. Basta adicionar quantos operadores AddCustom necessários. Além disso, BuiltinOpResolver também permite que você substitua 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 do TensorFlow Lite. Para fins de teste, você pode tornar sua versão local do TensorFlow Lite ciente de sua operação personalizada adicionando a chamada AddCustom apropriada (como mostrado acima) para 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 se deslocar (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 snippet 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 muito desperdício de memória, prefira usar um array estático 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 existem, pois 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 mantendo o tamanho binário pequeno. Veja o que outros kernels usam para obter informações (ou perguntar).

  5. Verifique o ponteiro para a memória retornada por malloc . Se esse ponteiro for nullptr , nenhuma operação deve 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 parada quando TF_LITE_ENSURE for usado, ou seja, essas macros devem ser usadas antes que qualquer recurso seja alocado que vaze.