Cette page a été traduite par l'API Cloud Translation.
Switch to English

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 autoriser 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 fusion des opérateurs .

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

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

Exemple: opérateur Sin personnalisé

Passons en revue un exemple de prise en charge d'un opérateur TensorFlow que TensorFlow Lite n'a pas. Supposons que nous utilisions l'opérateur Sin et que nous construisions un modèle très simple pour une fonction y = sin(x + offset) , où offset est entraînable.

Créer un modèle TensorFlow

L'extrait de code suivant entraîne un modèle TensorFlow simple. Ce modèle contient juste un opérateur personnalisé nommé Sin , qui est une fonction y = sin(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 = [-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

À 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:
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.

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

À ce stade, si vous l'exécutez avec l'interpréteur par défaut, vous obtiendrez les messages d'erreur suivants:

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

Créez et enregistrez l'opérateur.

Tous les opérateurs TensorFlow Lite (à la fois personnalisés et intégrés) sont définis à l'aide d'une interface pure-C simple 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'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() sera appelé plus d'une fois si l'op est utilisé plusieurs fois dans le graphique. Pour les opérations personnalisées, un tampon de configuration sera fourni, contenant un tampon flexible 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 op. 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 se débarrasser 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 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 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 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 séparées.

Définition du 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 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;
}

Lors de l'initialisation de l' OpResolver , ajoutez l'opération personnalisée dans le résolveur (voir ci-dessous 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 SinPrepare et SinEval vous avez définies pour l'opération personnalisée. Si vous SinInit SinFree fonctions SinInit et SinFree pour initialiser les variables utilisées dans l'opération et pour libérer de l'espace, 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 sera assignée pour exécuter chacun des opérateurs du modèle. Alors que la bibliothèque par défaut ne contient que des noyaux intégrés, il est possible de la remplacer / augmenter avec des opérateurs d'opérations de bibliothèque personnalisés.

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

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 BuiltinOpResolver et que vous BuiltinOpResolver :

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("Sin", Register_SIN());

Si l'ensemble des opérations intégrées est jugé trop grand, un nouveau OpResolver pourrait être généré par le 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 des 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 vous le souhaitez. En outre, BuiltinOpResolver vous permet également de remplacer les implémentations de fonctions intégrées à l'aide de 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 rendre votre version locale de TensorFlow Lite consciente 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 de mémoire et les désallocations 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 que dans chaque itération. Utilisez des données de tenseurs temporaires plutôt que de vous massacrer (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 de taille fixe statique (ou un std::vector Resize 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 la taille binaire petite. Découvrez 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 que vous avez une erreur de sortie, désallouez 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 que des ressources ne soient allouées qui fuiront.