RSVP para seu evento TensorFlow Everywhere local hoje!
Esta página foi traduzida pela API Cloud Translation.
Switch to English

Operadores personalizados

Como a biblioteca de operadores integrada do TensorFlow Lite oferece suporte a apenas um número limitado de operadores do TensorFlow, nem todo modelo é conversível. 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ê deseja combinar uma série de operadores TensorFlow sem suporte (ou com suporte) em um único operador personalizado otimizado fundido, consulte fusão do operador .

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 (denominado Sin , consulte #create_a_tensorflow_model) que é compatível com o TensorFlow, mas não é compatível com o TensorFlow Lite.

Exemplo: operador Custom Sin

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

Crie 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 denominado 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 as sinalizações do conversor padrão, 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 TensorFlow Lite

Crie um modelo TensorFlow Lite com operadores personalizados, definindo o atributo converter allow_custom_ops conforme mostrado abaixo:

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

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

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

Crie e registre o operador.

Todos os operadores 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. 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 gráfico. Um determinado init() será chamado mais de uma vez se o op for usado várias vezes no gráfico. Para operações customizadas, um buffer de configuração será fornecido, contendo um flexbuffer que mapeia nomes de 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 do 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 são redimensionados, o interpretador irá percorrer o gráfico notificando as implementações sobre a 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 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 atravessa 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 integradas, definindo essas quatro funções e uma função de registro global que geralmente se parece com isto:

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 dos builtins, as operações personalizadas terão que ser coletadas em bibliotecas personalizadas separadas.

Definição do kernel no tempo de execução 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 um exemplo abaixo). Isso registrará o operador com o Tensorflow Lite para que o TensorFlow Lite possa usar a nova implementação. Note-se que os dois últimos argumentos em TfLiteRegistration correspondem aos SinPrepare e SinEval funções definido para o op personalizado. Se você usou SinInit e SinFree funções para inicializar variáveis utilizadas na op e para liberar espaço, respectivamente, em seguida, eles seriam acrescentados 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 na biblioteca do kernel. Isso é feito com um OpResolver . Nos bastidores, o interpretador carregará uma biblioteca de kernels que será designada para executar cada um dos operadores no modelo. Embora a biblioteca padrão contenha apenas kernels embutidos, é possível substituí-la / aumentá-la com operadores op de biblioteca personalizados.

A classe OpResolver , que traduz códigos de operador e nomes 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 AddOp personalizada criada acima, chame AddOp (antes de passar o resolvedor para o InterpreterBuilder ):

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

Se o conjunto de ops embutidos for considerado muito grande, um novo OpResolver pode ser gerado por código com base em um determinado subconjunto de ops, possivelmente apenas aqueles contidos 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, atualmente precisa construir 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 Python .

Observe que um processo semelhante ao acima pode ser seguido para apoiar 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 de sua operadora

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 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 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 se maltratar (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. Você pode precisar usar a estrutura OpData para fazer referência aos índices de tensor em outras funções. Veja o exemplo no kernel para convolução . Um exemplo de snippet de código 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 um array estático de tamanho fixo (ou um std::vector pré-alocado em Resize ) ao invés 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, porque eles afetam o tamanho do binário. Por exemplo, se você precisa 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 binário pequeno. Veja o que outros kernels usam para obter uma visão (ou pergunte).

  5. Verifique o ponteiro para a memória retornada por malloc . Se este 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 a memória TF_LITE_ENSURE quando TF_LITE_ENSURE for usado, ou seja, essas macros devem ser usadas antes que quaisquer recursos que possam vazar sejam alocados.