Fusion d'opérations TensorFlow

Aperçu

Cette page décrit la conception et les étapes nécessaires pour convertir les opérations composites dans TensorFlow en opérations fusionnées dans TensorFlow Lite. Cette infrastructure est à usage général et prend en charge la conversion de toute opération composite dans TensorFlow en une opération fusionnée correspondante dans TensorFlow Lite.

Un exemple d'utilisation de cette infrastructure est la fusion des opérations TensorFlow RNN avec TensorFlow Lite, comme détaillé ici .

Que sont les opérations fusionnées

dessin

Les opérations TensorFlow peuvent être soit des opérations primitives, par exemple tf.add , soit elles peuvent être composées à partir d'autres opérations primitives, par exemple tf.einsum . Une opération primitive apparaît comme un nœud unique dans le graphique TensorFlow tandis qu'une opération composite est une collection de nœuds dans le graphique TensorFlow. Exécuter une opération composite équivaut à exécuter chacune de ses opérations primitives constitutives.

Une opération fusionnée correspond à une opération unique qui englobe tous les calculs effectués par chaque opération primitive au sein de l'opération composite correspondante.

Avantages des opérations fusionnées

Les opérations fusionnées existent pour maximiser les performances de leurs implémentations de noyau sous-jacentes, en optimisant le calcul global et en réduisant l'empreinte mémoire. Ceci est très utile, en particulier pour les charges de travail d'inférence à faible latence et les plates-formes mobiles aux ressources limitées.

Les opérations fusionnées fournissent également une interface de niveau supérieur pour définir des transformations complexes telles que la quantification, qui seraient autrement irréalisables ou très difficiles à réaliser à un niveau plus granulaire.

TensorFlow Lite comporte de nombreuses instances d'opérations fusionnées pour les raisons expliquées ci-dessus. Ces opérations fusionnées correspondent généralement aux opérations composites dans le programme TensorFlow source. Des exemples d'opérations composites dans TensorFlow qui sont implémentées en tant qu'opération fusionnée unique dans TensorFlow Lite incluent diverses opérations RNN telles que la séquence unidirectionnelle et bidirectionnelle LSTM, la convolution (conv2d, biais add, relu), entièrement connectée (matmul, biais add, relu) et plus encore. . Dans TensorFlow Lite, la quantification LSTM n'est actuellement implémentée que dans les opérations LSTM fusionnées.

Défis liés aux opérations fusionnées

La conversion d'opérations composites de TensorFlow en opérations fusionnées dans TensorFlow Lite est un problème difficile. Ceci est dû au fait:

  1. Les opérations composites sont représentées dans le graphique TensorFlow comme un ensemble d'opérations primitives sans limite bien définie. Il peut être très difficile d'identifier (par exemple via une correspondance de modèles) le sous-graphe correspondant à une telle opération composite.

  2. Il peut y avoir plusieurs implémentations TensorFlow ciblant une opération TensorFlow Lite fusionnée. Par exemple, il existe de nombreuses implémentations LSTM dans TensorFlow (Keras, Babelfish/lingvo, etc.) et chacune d'entre elles est composée de différentes opérations primitives, mais elles pourraient toutes être converties en la même opération LSTM fusionnée dans TensorFlow Lite.

En tant que telle, la conversion des opérations fusionnées s’est avérée assez difficile.

Enveloppez l'opération composite dans une tf.function

Dans de nombreux cas, une partie du modèle peut être mappée à une seule opération dans TFLite. Cela peut améliorer les performances lors de l’écriture d’une implémentation optimisée pour des opérations spécifiques. Pour pouvoir créer une opération fusionnée dans TFLite, identifiez la partie du graphique qui représente une opération fusionnée et enveloppez-la dans un tf.function avec l'attribut "experimental_implements" dans un tf.function , qui a la valeur d'attribut tfl_fusable_op avec la valeur true . Si l'opération personnalisée prend des attributs, transmettez-les dans le cadre du même "experimental_implements".

Exemple,

def get_implements_signature():
  implements_signature = [
    # 'name' will be used as a name for the operation.
    'name: "my_custom_fused_op"',
    # attr "tfl_fusable_op" is required to be set with true value.
    'attr {key: "tfl_fusable_op" value { b: true } }',
    # Example attribute "example_option" that the op accepts.
    'attr {key: "example_option" value { i: %d } }' % 10
  ]
  return ' '.join(implements_signature)

@tf.function(experimental_implements=get_implements_signature())
def my_custom_fused_op(input_1, input_2):
  # An empty function that represents pre/post processing example that
  # is not represented as part of the Tensorflow graph.
  output_1 = tf.constant(0.0, dtype=tf.float32, name='first_output')
  output_2 = tf.constant(0.0, dtype=tf.float32, name='second_output')
  return output_1, output_2

class TestModel(tf.Module):
  def __init__(self):
    super(TestModel, self).__init__()
    self.conv_1 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
    self.conv_2 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))

  @tf.function(input_signature=[
      tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
      tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
  ])
  def simple_eval(self, input_a, input_b):
    return my_custom_fused_op(self.conv_1(input_a), self.conv_2(input_b))

Notez que vous n'avez pas besoin de définir allow_custom_ops sur le convertisseur car l'attribut tfl_fusable_op l'implique déjà.

Implémentez une opération personnalisée et inscrivez-vous auprès de TFLite Interpreter

Implémentez votre opération fusionnée en tant qu'opération personnalisée TFLite - voir les instructions .

Notez que le nom avec lequel enregistrer l'opération doit être similaire au nom spécifié dans l'attribut name dans la signature d'implémentation.

Un exemple pour l'opération dans l'exemple est

  TfLiteRegistration reg = {};
  // This name must match the name specified in the implements signature.
  static constexpr char kOpName[] = "my_custom_fused_op";
  reg.custom_name = kOpName;
  reg.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
    // Add your code.
    return kTfLiteOk;
  };
  reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
    // Add your code.
    return kTfLiteOk;
  };
  reg.builtin_code = kTfLiteCustom;
  resolver->AddCustom(kOpName, &reg);

Conversion d'un fonctionnement composite à un fonctionnement fusionné (avancé)

L'architecture globale de conversion des opérations composites TensorFlow en opérations fusionnées TensorFlow Lite est la suivante :

dessin

Enveloppez l'opération composite dans une tf.function

Dans le code source du modèle TensorFlow, identifiez et extrayez l'opération composite dans un tf.function avec l'annotation de fonction experimental_implements . Voir un exemple de recherche intégrée . La fonction définit l'interface et ses arguments doivent être utilisés pour implémenter la logique de conversion.

Écrire le code de conversion

Le code de conversion est écrit par l'interface de la fonction avec l'annotation implements . Voir un exemple de fusion pour intégrer la recherche . Conceptuellement, le code de conversion remplace l'implémentation composite de cette interface par celle fusionnée.

Dans la passe prepare-composite-functions, insérez votre code de conversion .

Dans des usages plus avancés, il est possible d'implémenter des transformations complexes des opérandes de l'opération composite afin de dériver les opérandes de l'opération fusionnée. Voir Keras LSTM . code de conversion à titre d'exemple.

Convertir en TensorFlow Lite

Utilisez l'API TFLiteConverter.from_saved_model pour convertir en TensorFlow Lite.

Sous la capuche

Nous décrivons maintenant les détails de haut niveau de la conception globale de la conversion en opérations fusionnées dans TensorFlow Lite.

Composer des opérations dans TensorFlow

L'utilisation de tf.function avec l'attribut de fonction experimental_implements permet aux utilisateurs de composer explicitement de nouvelles opérations à l'aide d'opérations primitives TensorFlow et de spécifier l'interface implémentée par l'opération composite résultante. Ceci est très utile car il fournit :

  1. Une limite bien définie pour l'opération composite dans le graphique TensorFlow sous-jacent.
  2. Spécifiez explicitement l’interface implémentée par cette opération. Les arguments de la tf.function correspondent aux arguments de cette interface.

À titre d'exemple, considérons une opération composite définie pour implémenter la recherche d'intégration. Cela correspond à une opération fusionnée dans TensorFlow Lite.

  @tf.function(
        experimental_implements="embedding_lookup")
    def EmbFprop(embs, ids_vec):
      """Embedding forward prop.

      Effectively, it computes:
        num = size of ids_vec
        rets = zeros([num, embedding dim])
        for i in range(num):
          rets[i, :] = embs[ids_vec[i], :]
        return rets

      Args:
        embs: The embedding matrix.
        ids_vec: A vector of int32 embedding ids.

      Returns:
        The result of embedding lookups. A matrix of shape
        [num ids in ids_vec, embedding dims].
      """
      num = tf.shape(ids_vec)[0]
      rets = inplace_ops.empty([num] + emb_shape_suf, py_utils.FPropDtype(p))

      def EmbFpropLoop(i, embs, ids_vec, rets):
        # row_id = ids_vec[i]
        row_id = tf.gather(ids_vec, i)
        # row = embs[row_id]
        row = tf.reshape(tf.gather(embs, row_id), [1] + emb_shape_suf)
        # rets[i] = row
        rets = inplace_ops.alias_inplace_update(rets, [i], row)
        return embs, ids_vec, rets

      _, _, rets = functional_ops.For(
          start=0,
          limit=num,
          delta=1,
          inputs=[embs, ids_vec, rets],
          body=EmbFpropLoop,
          rewrite_with_while=compiled)
      if len(weight_shape) > 2:
        rets = tf.reshape(rets, [num, symbolic.ToStatic(p.embedding_dim)])
      return rets

En faisant en sorte que les modèles utilisent des opérations composites via tf.function comme illustré ci-dessus, il devient possible de créer une infrastructure générale pour identifier et convertir ces opérations en opérations TensorFlow Lite fusionnées.

Extension du convertisseur TensorFlow Lite

Le convertisseur TensorFlow Lite publié plus tôt cette année ne prenait en charge que l'importation de modèles TensorFlow sous forme de graphique avec toutes les variables remplacées par leurs valeurs constantes correspondantes. Cela ne fonctionne pas pour l'opération de fusion puisque ces graphiques ont toutes les fonctions intégrées afin que les variables puissent être transformées en constantes.

Afin d'exploiter tf.function avec la fonctionnalité experimental_implements pendant le processus de conversion, les fonctions doivent être conservées jusqu'à plus tard dans le processus de conversion.

En tant que tel, nous avons mis en œuvre un nouveau flux de travail d'importation et de conversion de modèles TensorFlow dans le convertisseur pour prendre en charge le cas d'utilisation de la fusion d'opérations composites. Concrètement, les nouvelles fonctionnalités ajoutées sont :

  1. Importer des modèles enregistrés TensorFlow dans MLIR
  2. fusionner les opérations composites
  3. analyse de mutabilité variable
  4. geler toutes les variables en lecture seule

Cela nous permet d'effectuer une fusion d'opérations en utilisant les fonctions représentant les opérations composites avant l'intégration des fonctions et le gel des variables.

Implémentation de l'opération fusion

Examinons plus en détail l'opération de fusion. Ce pass effectue les opérations suivantes :

  1. Parcourez toutes les fonctions du module MLIR.
  2. Si une fonction possède l'attribut tf._implements, en fonction de la valeur de l'attribut, elle appelle l'utilitaire de fusion d'opérations approprié.
  3. L'utilitaire de fusion d'opérations opère sur les opérandes et les attributs de la fonction (qui servent d'interface pour la conversion) et remplace le corps de la fonction par un corps de fonction équivalent contenant l'opération fusionnée.
  4. Dans de nombreux cas, le corps remplacé contiendra des opérations autres que l’opération fusionnée. Celles-ci correspondent à des transformations statiques sur les opérandes de la fonction afin d'obtenir les opérandes de l'opération fusionnée. Étant donné que ces calculs peuvent tous être repliés de manière constante, ils ne seraient pas présents dans le tampon plat exporté où seule l'opération fusionnée existerait.

Voici un extrait de code de la passe montrant le flux de travail principal :

void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
                                                        StringAttr attr) {
  if (attr.getValue() == "embedding_lookup") {
    func.eraseBody();
    func.addEntryBlock();
    // Convert the composite embedding_lookup function body to a
    // TFLite fused embedding_lookup op.
    ConvertEmbeddedLookupFunc convert_embedded_lookup(func);
    if (failed(convert_embedded_lookup.VerifySignature())) {
      return signalPassFailure();
    }
    convert_embedded_lookup.RewriteFunc();
  } else if (attr.getValue() == mlir::TFL::kKerasLstm) {
     func.eraseBody();
     func.addEntryBlock();
     OpBuilder builder(func.getBody());
     if (failed(ConvertKerasLSTMLayer(func, &builder))) {
       return signalPassFailure();
     }
  } else if (.....) /* Other fusions can plug in here */
}

Voici un extrait de code montrant le mappage de cette opération composite à une opération fusionnée dans TensorFlow Lite en exploitant la fonction comme interface de conversion.

void RewriteFunc() {
    Value lookup = func_.getArgument(1);
    Value value = func_.getArgument(0);
    auto output_type = func_.getType().getResult(0);

    OpBuilder builder(func_.getBody());
    auto op = builder.create<mlir::TFL::EmbeddingLookupOp>(
        func_.getLoc(), output_type, lookup, value);

    builder.create<mlir::ReturnOp>(func_.getLoc(), op.getResult());
  }