Operadores personalizados

Dado que la biblioteca de operadores integrada de TensorFlow Lite solo admite una cantidad limitada de operadores de TensorFlow, no todos los modelos son convertibles. Para obtener más información, consulte la compatibilidad del operador .

Para permitir la conversión, los usuarios pueden proporcionar su propia implementación personalizada de un operador TensorFlow no compatible en TensorFlow Lite, conocido como operador personalizado. Si, en cambio, desea combinar una serie de operadores de TensorFlow no admitidos (o admitidos) en un solo operador personalizado optimizado fusionado, consulte la fusión de operadores .

El uso de operadores personalizados consta de cuatro pasos.

Veamos un ejemplo completo de ejecución de un modelo con un operador personalizado tf.atan (llamado Atan , consulte #create_a_tensorflow_model) que es compatible con TensorFlow, pero no con TensorFlow Lite.

El operador de texto de TensorFlow es un ejemplo de un operador personalizado. Consulte el tutorial Convertir texto TF a TF Lite para ver un ejemplo de código.

Ejemplo: operador Atan personalizado

Veamos un ejemplo de compatibilidad con un operador de TensorFlow que TensorFlow Lite no tiene. Supongamos que estamos usando el operador Atan y que estamos construyendo un modelo muy simple para una función y = atan(x + offset) , donde offset es entrenable.

Crear un modelo de TensorFlow

El siguiente fragmento de código entrena un modelo simple de TensorFlow. Este modelo solo contiene un operador personalizado llamado Atan , que es una función y = atan(x + offset) , donde offset se puede entrenar.

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

En este punto, si intenta generar un modelo de TensorFlow Lite con las banderas de convertidor predeterminadas, obtendrá el siguiente mensaje de error:

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

Convertir a un modelo TensorFlow Lite

Cree un modelo de TensorFlow Lite con operadores personalizados configurando el atributo del convertidor allow_custom_ops como se muestra a continuación:

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

En este punto, si lo ejecuta con el intérprete predeterminado usando comandos como los siguientes:

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

Seguirá recibiendo el error:

Encountered unresolved custom op: Atan.

Crear y registrar el operador.

Todos los operadores de TensorFlow Lite (tanto personalizados como integrados) se definen mediante una interfaz simple de C puro que consta de cuatro funciones:

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 obtener detalles sobre TfLiteContext y TfLiteNode . El primero proporciona funciones de informe de errores y acceso a objetos globales, incluidos todos los tensores. Este último permite que las implementaciones accedan a sus entradas y salidas.

Cuando el intérprete carga un modelo, llama a init() una vez por cada nodo del gráfico. Se llamará a un init() dado más de una vez si la operación se usa varias veces en el gráfico. Para operaciones personalizadas, se proporcionará un búfer de configuración que contiene un búfer flexible que asigna los nombres de los parámetros a sus valores. El búfer está vacío para operaciones integradas porque el intérprete ya analizó los parámetros de operación. Las implementaciones de kernel que requieren estado deben inicializarlo aquí y transferir la propiedad a la persona que llama. Para cada llamada init() , habrá una llamada correspondiente a free() , lo que permitirá que las implementaciones eliminen el búfer que podrían haber asignado en init() .

Siempre que se cambie el tamaño de los tensores de entrada, el intérprete pasará por el gráfico notificando las implementaciones del cambio. Esto les da la oportunidad de cambiar el tamaño de su búfer interno, verificar la validez de las formas y tipos de entrada y volver a calcular las formas de salida. Todo esto se hace a través de prepare() , y las implementaciones pueden acceder a su estado usando node->user_data .

Finalmente, cada vez que se ejecuta la inferencia, el intérprete recorre el gráfico llamando invoke() , y aquí también el estado está disponible como node->user_data .

Las operaciones personalizadas se pueden implementar exactamente de la misma manera que las operaciones integradas, al definir esas cuatro funciones y una función de registro global que generalmente se ve así:

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

Tenga en cuenta que el registro no es automático y se debe realizar una llamada explícita a Register_MY_CUSTOM_OP . Si bien el BuiltinOpResolver estándar (disponible en el destino :builtin_ops ) se encarga del registro de las funciones integradas, las operaciones personalizadas deberán recopilarse en bibliotecas personalizadas separadas.

Definición del kernel en el tiempo de ejecución de TensorFlow Lite

Todo lo que tenemos que hacer para usar el op en TensorFlow Lite es definir dos funciones ( Prepare y Eval ) y construir un 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;
}

Al inicializar el OpResolver , agregue la operación personalizada en el resolver (vea un ejemplo a continuación). Esto registrará al operador con Tensorflow Lite para que TensorFlow Lite pueda usar la nueva implementación. Tenga en cuenta que los dos últimos argumentos en TfLiteRegistration corresponden a las funciones AtanPrepare y AtanEval que definió para la operación personalizada. Si utilizó las funciones AtanInit y AtanFree para inicializar las variables utilizadas en la operación y para liberar espacio, respectivamente, se agregarían a los dos primeros argumentos de TfLiteRegistration ; esos argumentos se establecen en nullptr en este ejemplo.

Registrar el operador con la biblioteca del kernel

Ahora necesitamos registrar el operador con la biblioteca del núcleo. Esto se hace con un OpResolver . Detrás de escena, el intérprete cargará una biblioteca de kernels que se asignarán para ejecutar cada uno de los operadores en el modelo. Si bien la biblioteca predeterminada solo contiene núcleos integrados, es posible reemplazarla/aumentarla con operadores operativos de biblioteca personalizados.

La clase OpResolver , que traduce los códigos y nombres de operadores en código real, se define así:

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

El uso regular requiere que use BuiltinOpResolver y escriba:

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

Para agregar la operación personalizada creada anteriormente, llame a AddOp (antes de pasar la resolución al InterpreterBuilder ):

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

Si se considera que el conjunto de operaciones integradas es demasiado grande, se podría generar un código OpResolver nuevo en función de un subconjunto determinado de operaciones, posiblemente solo las contenidas en un modelo determinado. Este es el equivalente al registro selectivo de TensorFlow (y una versión simple está disponible en el directorio tools ).

Si desea definir sus operadores personalizados en Java, actualmente necesitaría crear su propia capa JNI personalizada y compilar su propio AAR en este código jni . Del mismo modo, si desea definir estos operadores disponibles en Python, puede colocar sus registros en el código contenedor de Python .

Tenga en cuenta que se puede seguir un proceso similar al anterior para admitir un conjunto de operaciones en lugar de un solo operador. Simplemente agregue tantos operadores AddCustom como necesite. Además, BuiltinOpResolver también le permite anular implementaciones de elementos integrados mediante AddBuiltin .

Pruebe y perfile a su operador

Para perfilar su operación con la herramienta de evaluación comparativa TensorFlow Lite, puede usar la herramienta de modelo de evaluación comparativa para TensorFlow Lite. Con fines de prueba, puede hacer que su compilación local de TensorFlow Lite conozca su operación personalizada agregando la llamada AddCustom adecuada (como se muestra arriba) a register.cc

Mejores prácticas

  1. Optimice las asignaciones y desasignaciones de memoria con precaución. La asignación de memoria en Prepare es más eficiente que en Invoke y la asignación de memoria antes de un ciclo es mejor que en cada iteración. Utilice datos de tensores temporales en lugar de mallocarse usted mismo (consulte el punto 2). Use punteros/referencias en lugar de copiar tanto como sea posible.

  2. Si una estructura de datos persistirá durante toda la operación, recomendamos preasignar la memoria utilizando tensores temporales. Es posible que deba usar la estructura OpData para hacer referencia a los índices de tensor en otras funciones. Vea el ejemplo en el kernel para convolución . Un fragmento de código de muestra está debajo

    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. Si no cuesta demasiada memoria desperdiciada, prefiera usar una matriz estática de tamaño fijo (o un std::vector preasignado en Resize ) en lugar de usar un std::vector asignado dinámicamente en cada iteración de ejecución.

  4. Evite crear instancias de plantillas de contenedor de biblioteca estándar que aún no existen, ya que afectan el tamaño binario. Por ejemplo, si necesita un std::map en su operación que no existe en otros kernels, el uso de un std::vector con mapeo de indexación directa podría funcionar manteniendo el tamaño binario pequeño. Vea lo que usan otros núcleos para obtener información (o preguntar).

  5. Verifique el puntero a la memoria devuelta por malloc . Si este puntero es nullptr , no se deben realizar operaciones con ese puntero. Si malloc en una función y tiene una salida de error, desasigne la memoria antes de salir.

  6. Use TF_LITE_ENSURE(context, condition) para verificar una condición específica. Su código no debe dejar la memoria colgando cuando se usa TF_LITE_ENSURE , es decir, estas macros deben usarse antes de que se asignen recursos que se filtren.