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 à 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 à fusion d'opérateurs .

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 TF Text 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 ne possède 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 entraîné.

Créer un modèle TensorFlow

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

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 un modèle TensorFlow Lite

Créez un modèle TensorFlow Lite avec des opérateurs personnalisés en définissant l'attribut de 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 en utilisant des commandes telles que celles-ci :

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 fonctionnalités de rapport d’erreurs 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 graphique. Un init() donné sera appelé plus d'une fois si l'opération est utilisée plusieurs fois dans le graphique. 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 de 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 parcourt le graphique pour notifier 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 sous la forme 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 my_namespace {
  const TfLiteRegistration* Register_MY_CUSTOM_OP() {
    static const TfLiteRegistration r = {my_custom_op::Init,
                                         my_custom_op::Free,
                                         my_custom_op::Prepare,
                                         my_custom_op::Eval};
    return &r;
  }
}  // namespace my_namespace

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 fonctions intégrées, les opérations personnalisées devront être collectées dans des bibliothèques personnalisées distinctes.

Définir le noyau dans le runtime 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;
}

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

Lors de l'initialisation de OpResolver , ajoutez l'opération personnalisée 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 de TfLiteRegistration correspondent aux fonctions AtanPrepare et AtanEval que vous avez définies pour l'opération personnalisée. Si vous utilisiez les fonctions AtanInit et AtanFree pour initialiser les variables utilisées dans l'opération 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 auprès de la bibliothèque du noyau

Nous devons maintenant enregistrer l'opérateur auprès de la bibliothèque du noyau. Cela se fait avec un OpResolver . En 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 par une bibliothèque d'opérateurs personnalisée.

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

class OpResolver {
 public:
  virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
  virtual TfLiteRegistration* FindOp(const char* op) const = 0;
  ...
};

Les classes MutableOpResolver et BuiltinOpResolver sont dérivées de OpResolver :

class MutableOpResolver : public OpResolver {
 public:
  MutableOpResolver();  // Constructs an initially empty op resolver.
  void AddBuiltin(tflite::BuiltinOperator op, const TfLiteRegistration* registration) = 0;
  void AddCustom(const char* op, const TfLiteRegistration* registration) = 0;
  void AddAll(const MutableOpResolver& other);
  ...
};

class BuiltinOpResolver : public MutableOpResolver {
 public:
  BuiltinOpResolver();  // Constructs an op resolver with all the builtin ops.
};

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

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

Pour ajouter l'opération personnalisée créée ci-dessus, vous pouvez à la place utiliser un MutableOpResolver et appeler AddCustom (avant de transmettre le résolveur au InterpreterBuilder ) :

tflite::ops::builtin::MutableOpResolver resolver;
resolver.AddAll(tflite::ops::builtin::BuiltinOpResolver());
resolver.AddCustom("Atan", Register_ATAN());

Si l'ensemble des opérations intégrées est jugé trop volumineux, un nouvel OpResolver pourrait être généré par code sur la base d'un sous-ensemble d'opérations donné, é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 devrez 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, MutableOpResolver vous permet également de remplacer les implémentations des fonctions intégrées en utilisant 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 les désallocations de mémoire avec prudence. 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 tenseurs temporaires plutôt que de vous malallouer (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 tensoriels dans d'autres fonctions. Voir l'exemple dans le noyau pour la 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 conteneurs 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 mappage d'indexation direct pourrait fonctionner tout en gardant la taille binaire petite. Découvrez ce que les autres noyaux utilisent pour obtenir un aperçu (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 une fonction et que vous rencontrez une erreur de sortie, libérez 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 la 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.