Opérateurs personnalisés

Étant donné que la bibliothèque d'opérateurs intégrée TensorFlow Lite ne prend en charge qu'un nombre limité d'opérateurs TensorFlow, tous les modèles ne sont pas convertibles. Pour plus de détails, reportez-vous à la compatibilité des opérateurs .

Pour permettre la conversion, les utilisateurs peuvent fournir leur propre implémentation personnalisée d'un opérateur TensorFlow non pris en charge dans TensorFlow Lite, appelé opérateur personnalisé. Si, à la place, vous souhaitez combiner une série d'opérateurs TensorFlow non pris en charge (ou pris en charge) en un seul opérateur personnalisé optimisé fusionné, reportez-vous à la section operator fusing .

L'utilisation d'opérateurs personnalisés comprend quatre étapes.

Passons en revue un exemple de bout en bout d'exécution d'un modèle avec un opérateur personnalisé tf.atan (nommé Atan , reportez-vous à #create_a_tensorflow_model) qui est pris en charge dans TensorFlow, mais non pris en charge dans TensorFlow Lite.

L'opérateur TensorFlow Text est un exemple d'opérateur personnalisé. Consultez le didacticiel Convertir du texte TF en TF Lite pour un exemple de code.

Exemple : Opérateur Atan personnalisé

Passons en revue un exemple de prise en charge d'un opérateur TensorFlow que TensorFlow Lite n'a pas. Supposons que nous utilisons l'opérateur Atan et que nous construisons un modèle très simple pour une fonction y = atan(x + offset) , où offset peut être formé.

Créer un modèle TensorFlow

L'extrait de code suivant forme un modèle TensorFlow simple. Ce modèle contient juste un opérateur personnalisé nommé Atan , qui est une fonction y = atan(x + offset) , où offset peut être formé.

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

À ce stade, si vous essayez de générer un modèle TensorFlow Lite avec les indicateurs de convertisseur par défaut, vous obtiendrez le message d'erreur suivant :

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

Convertir en modèle TensorFlow Lite

Créez un modèle TensorFlow Lite avec des opérateurs personnalisés en définissant l'attribut du convertisseur allow_custom_ops comme indiqué ci-dessous :

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

À ce stade, si vous l'exécutez avec l'interpréteur par défaut à l'aide de commandes telles que :

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

Vous obtiendrez toujours l'erreur :

Encountered unresolved custom op: Atan.

Créez et enregistrez l'opérateur.

Tous les opérateurs TensorFlow Lite (personnalisés et intégrés) sont définis à l'aide d'une simple interface en C pur composée de quatre fonctions :

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;

Reportez-vous à common.h pour plus de détails sur TfLiteContext et TfLiteNode . Le premier fournit des fonctions de rapport d'erreur et un accès aux objets globaux, y compris tous les tenseurs. Ce dernier permet aux implémentations d'accéder à leurs entrées et sorties.

Lorsque l'interpréteur charge un modèle, il appelle init() une fois pour chaque nœud du graphe. Un init() donné sera appelé plus d'une fois si l'op est utilisé plusieurs fois dans le graphe. Pour les opérations personnalisées, un tampon de configuration sera fourni, contenant un flexbuffer qui mappe les noms de paramètres à leurs valeurs. Le tampon est vide pour les opérations intégrées car l'interpréteur a déjà analysé les paramètres de l'opération. Les implémentations du noyau qui nécessitent un état doivent l'initialiser ici et transférer la propriété à l'appelant. Pour chaque appel init() , il y aura un appel correspondant à free() , permettant aux implémentations de disposer du tampon qu'elles auraient pu allouer dans init() .

Chaque fois que les tenseurs d'entrée sont redimensionnés, l'interpréteur parcourra le graphique notifiant les implémentations du changement. Cela leur donne la possibilité de redimensionner leur tampon interne, de vérifier la validité des formes et des types d'entrée et de recalculer les formes de sortie. Tout cela se fait via prepare() , et les implémentations peuvent accéder à leur état en utilisant node->user_data .

Enfin, chaque fois que l'inférence s'exécute, l'interpréteur parcourt le graphe en appelant invoke() , et ici aussi l'état est disponible en tant que node->user_data .

Les opérations personnalisées peuvent être implémentées exactement de la même manière que les opérations intégrées, en définissant ces quatre fonctions et une fonction d'enregistrement globale qui ressemble généralement à ceci :

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

Notez que l'enregistrement n'est pas automatique et qu'un appel explicite à Register_MY_CUSTOM_OP doit être effectué. Alors que le BuiltinOpResolver standard (disponible à partir de la cible :builtin_ops ) s'occupe de l'enregistrement des builtins, les opérations personnalisées devront être collectées dans des bibliothèques personnalisées distinctes.

Définir le noyau dans l'environnement d'exécution TensorFlow Lite

Tout ce que nous devons faire pour utiliser l'op dans TensorFlow Lite est de définir deux fonctions ( Prepare et Eval ) et de construire 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;
}

Lors de l'initialisation de OpResolver , ajoutez l'op personnalisé dans le résolveur (voir ci-dessous pour un exemple). Cela enregistrera l'opérateur auprès de Tensorflow Lite afin que TensorFlow Lite puisse utiliser la nouvelle implémentation. Notez que les deux derniers arguments dans TfLiteRegistration correspondent aux fonctions AtanPrepare et AtanEval que vous avez définies pour l'op personnalisé. Si vous utilisiez les fonctions AtanInit et AtanFree pour initialiser les variables utilisées dans l'op et pour libérer de l'espace, respectivement, elles seraient alors ajoutées aux deux premiers arguments de TfLiteRegistration ; ces arguments sont définis sur nullptr dans cet exemple.

Enregistrez l'opérateur avec la bibliothèque du noyau

Nous devons maintenant enregistrer l'opérateur avec la bibliothèque du noyau. Ceci est fait avec un OpResolver . Dans les coulisses, l'interpréteur chargera une bibliothèque de noyaux qui seront affectés à l'exécution de chacun des opérateurs du modèle. Bien que la bibliothèque par défaut ne contienne que des noyaux intégrés, il est possible de la remplacer/l'augmenter avec des opérateurs op de bibliothèque personnalisés.

La classe OpResolver , qui traduit les codes et les noms des opérateurs en code réel, est définie comme suit :

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

Une utilisation régulière nécessite que vous utilisiez BuiltinOpResolver et que vous écriviez :

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

Pour ajouter l'op personnalisé créé ci-dessus, vous appelez AddOp (avant de passer le résolveur à InterpreterBuilder ):

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

Si l'ensemble d'opérations intégrées est jugé trop volumineux, un nouvel OpResolver peut être généré par code en fonction d'un sous-ensemble donné d'opérations, éventuellement uniquement celles contenues dans un modèle donné. C'est l'équivalent de l'enregistrement sélectif de TensorFlow (et une version simple de celui-ci est disponible dans le répertoire tools ).

Si vous souhaitez définir vos opérateurs personnalisés en Java, vous devez actuellement créer votre propre couche JNI personnalisée et compiler votre propre AAR dans ce code jni . De même, si vous souhaitez définir ces opérateurs disponibles en Python vous pouvez placer vos inscriptions dans le code wrapper Python .

Notez qu'un processus similaire à celui ci-dessus peut être suivi pour prendre en charge un ensemble d'opérations au lieu d'un seul opérateur. Ajoutez simplement autant d'opérateurs AddCustom que nécessaire. De plus, BuiltinOpResolver vous permet également de remplacer les implémentations de builtins en utilisant le AddBuiltin .

Testez et profilez votre opérateur

Pour profiler votre opération avec l'outil de référence TensorFlow Lite, vous pouvez utiliser l' outil de modèle de référence pour TensorFlow Lite. À des fins de test, vous pouvez informer votre version locale de TensorFlow Lite de votre opération personnalisée en ajoutant l'appel AddCustom approprié (comme indiqué ci-dessus) à register.cc

Les meilleures pratiques

  1. Optimisez les allocations et désallocations de mémoire avec précaution. L'allocation de mémoire dans Prepare est plus efficace que dans Invoke , et l'allocation de mémoire avant une boucle est meilleure qu'à chaque itération. Utilisez des données de tenseurs temporaires plutôt que de mallocer vous-même (voir point 2). Utilisez des pointeurs/références au lieu de copier autant que possible.

  2. Si une structure de données persiste pendant toute l'opération, nous vous conseillons de pré-allouer la mémoire à l'aide de tenseurs temporaires. Vous devrez peut-être utiliser la structure OpData pour référencer les indices de tenseur dans d'autres fonctions. Voir l'exemple dans le noyau pour convolution . Un exemple d'extrait de code est ci-dessous

    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 cela ne coûte pas trop de mémoire gaspillée, préférez utiliser un tableau statique de taille fixe (ou un std::vector pré-alloué dans Resize ) plutôt que d'utiliser un std::vector alloué dynamiquement à chaque itération d'exécution.

  4. Évitez d'instancier des modèles de conteneur de bibliothèque standard qui n'existent pas déjà, car ils affectent la taille binaire. Par exemple, si vous avez besoin d'un std::map dans votre opération qui n'existe pas dans d'autres noyaux, l'utilisation d'un std::vector avec un mappage d'indexation directe pourrait fonctionner tout en gardant une petite taille binaire. Voyez ce que les autres noyaux utilisent pour obtenir des informations (ou demander).

  5. Vérifiez le pointeur vers la mémoire renvoyée par malloc . Si ce pointeur est nullptr , aucune opération ne doit être effectuée à l'aide de ce pointeur. Si vous malloc dans une fonction et avez une sortie d'erreur, libérez de la mémoire avant de quitter.

  6. Utilisez TF_LITE_ENSURE(context, condition) pour vérifier une condition spécifique. Votre code ne doit pas laisser de mémoire en suspens lorsque TF_LITE_ENSURE est utilisé, c'est-à-dire que ces macros doivent être utilisées avant l'allocation de ressources susceptibles de fuir.