TensorFlow Operations Fusion

Présentation

Cette page décrit la conception et les étapes requises pour convertir des opérations composites de TensorFlow aux opérations fusionnées dans LiteRT. Cette infrastructure et permet la conversion de toute opération composite dans TensorFlow à une opération fusionnée correspondante dans LiteRT.

Un exemple d'utilisation de cette infrastructure est la fusion d'opérations TensorFlow RNN pour LiteRT, comme indiqué sur cette page.

Que sont les opérations fusionnées ?

dessin

Les opérations TensorFlow peuvent être des opérations primitives, par exemple tf.add ou peuvent être composé d'autres opérations primitives, comme tf.einsum. Une primitive apparaît comme un nœud unique dans le graphe TensorFlow, tandis qu'un nœud composite est un ensemble de nœuds dans le graphe TensorFlow. Exécuter une opération composite équivaut à exécuter chacune des primitives qui la composent. opérations.

Une opération fusionnée correspond à une opération unique qui regroupe toutes les de calcul effectué par chaque opération primitive dans le champ opération composite.

Avantages des opérations fusionnées

Les opérations fusionnées permettent d'optimiser les performances du noyau sous-jacent en optimisant le calcul global et en réduisant la mémoire de sécurité. C'est très utile, en particulier pour les charges de travail d'inférence à faible latence et les plates-formes mobiles à ressources limitées.

Les opérations fusionnées fournissent également une interface de niveau supérieur permettant de définir des des transformations comme la quantification, ce qui serait sinon impossible difficile à faire à un niveau plus précis.

LiteRT propose de nombreuses instances de fusions pour les raisons suivantes : articulés ci-dessus. Ces opérations fusionnées correspondent généralement à des dans le programme TensorFlow source. Exemples d'opérations composites dans TensorFlow qui sont implémentés en tant qu'opération fusionnée unique dans LiteRT. incluent diverses opérations RNN telles que les séquences unidirectionnelles et bidirectionnelles LSTM, convolution (conv2d, bias add, relu), entièrement connecté (matmul, biais add, relu) et plus encore. Dans LiteRT, la quantification LSTM n'est actuellement implémentées dans les opérations Fused LSTM.

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

Convertir des opérations composites de TensorFlow en opérations fusionnées dans Le format LiteRT est un problème difficile. Voici pourquoi :

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

  2. Plusieurs implémentations TensorFlow peuvent cibler une combinaison fusionnée Opération Lite. Par exemple, de nombreuses implémentations LSTM dans TensorFlow (Keras, Babelfish/Lingvo, etc.). Chacune d'entre elles opérations primitives différentes, mais elles peuvent toutes être converties en la même opération Fused LSTM dans LiteRT.

Ainsi, la conversion des opérations fusionnées s'est révélée assez difficile.

Encapsuler l'opération composite dans un tf.function

Dans de nombreux cas, une partie du modèle peut être mappée en 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, identifier la partie du graphique qui représente une opération fusionnée et l'encapsuler dans une fonction tf.function avec "experimental_implémentes" à une tf.function, qui possède l'attribut la valeur tfl_fusable_op avec la valeur true. Si l'opération personnalisée prend puis les transmettre 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 qu'il n'est pas nécessaire de définir allow_custom_ops sur le convertisseur comme L'attribut tfl_fusable_op implique déjà.

Implémenter des opérations et des enregistrements personnalisés avec l'interpréteur TFLite

Implémentez votre opération fusionnée en tant qu'opération personnalisée TFLite (voir la page). instructions.

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

Un exemple pour l'opération dans cet 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);

Convertir une opération composite en opération fusionnée (avancé)

L'architecture globale permettant de convertir des opérations composites TensorFlow Vous trouverez ci-dessous les opérations fusionnées en mode LiteRT:

dessin

Encapsuler l'opération composite dans un tf.function

Dans le code source du modèle TensorFlow, identifiez et extrayez le code composite dans une tf.function par le biais experimental_implements annotation de fonction. Consultez un exemple de recherche de représentation vectorielle continue. La définit l'interface et ses arguments doivent être utilisés pour implémenter la fonction logique de conversion.

Écrire le code de conversion

Le code de conversion est écrit conformément à l'interface de la fonction avec la variable Annotation implements. Voir un exemple de fusion pour les représentations vectorielles continues recherche. Sur le plan conceptuel, le code de conversion remplace l'ensemble de données l'implémentation de cette interface avec celle fusionnée.

À l'étape prepare-composite-functions, plug-in de votre conversion, code.

Dans les utilisations plus avancées, il est possible d'implémenter des transformations complexes de les opérandes de l'opération composite afin de dériver les opérandes de la fusion opération. Voir Keras LSTM. du code de conversion.

Convertir au format LiteRT

Utilisez le TFLiteConverter.from_saved_model pour la convertir au format LiteRT.

dans le détail

Nous décrivons maintenant les détails généraux de la conception globale lors de la conversion en dans LiteRT.

Opérations de composition dans TensorFlow

Utilisation de tf.function avec le experimental_implements L'attribut de fonction permet aux utilisateurs de composer explicitement de nouvelles opérations à l'aide de primitives de TensorFlow, et spécifie l'interface que le résultat implémentations d'opérations composites. Cette fonctionnalité est très utile, car elle fournit les éléments suivants:

  1. Une limite bien définie pour l'opération composite dans le Graphe TensorFlow.
  2. Spécifiez explicitement l'interface que cette opération implémente. La les arguments de la fonction tf.function correspondent aux arguments de cette interface.

Prenons l'exemple d'une opération composite définie pour implémenter de représentations vectorielles continues. Cela correspond à une opération fusionnée dans LiteRT.

  @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

Les modèles utilisent des opérations composites via tf.function en tant que illustré ci-dessus, il devient possible de créer une infrastructure générale identifier et convertir ces opérations en opérations LiteRT fusionnées.

Extension du convertisseur LiteRT

Le convertisseur LiteRT sorti plus tôt cette année n'est compatible qu'avec en important des modèles TensorFlow sous forme de graphe dans lequel toutes les variables sont remplacées par les valeurs constantes correspondantes. Cela ne fonctionne pas pour la fusion d'opérations ces graphiques ont toutes les fonctions intégrées afin que les variables puissent être transformées en constantes.

Afin de tirer parti tf.function par le biais experimental_implements pendant le processus de conversion, les fonctions doivent être conservées jusqu'à la fin du processus de conversion.

À ce titre, nous avons mis en place un nouveau workflow d'importation et de conversion de TensorFlow dans le convertisseur pour prendre en charge le cas d'utilisation de fusion d'opérations composites. Plus précisément, les nouvelles fonctionnalités sont les suivantes:

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

Cela nous permet d'effectuer une fusion d'opérations à l'aide des fonctions représentant des opérations composites avant l'intégration de la fonction et le gel des variables.

Implémenter la fusion des opérations

Examinons de plus près l'opération Fusion pass. Cette passe effectue suivantes:

  1. Parcourez toutes les fonctions du module MLIR.
  2. Si une fonction possède l'attribut tf._implements, en fonction de l'attribut appelle l'utilitaire de fusion d'opérations approprié.
  3. L'utilitaire de fusion d'opérations opère sur les opérandes de la fonction et (qui servent d'interface pour la conversion) et remplace le corps de la fonction par un corps de fonction équivalent contenant le opération fusionnée.
  4. Dans de nombreux cas, le corps remplacé contiendra des opérations autres que opération fusionnée. Elles correspondent à des transformations statiques 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 déployés en continu, ils ne sont pas présent dans le flatbuffer exporté, où seule l'opération fusionnée serait existent.

Voici l'extrait de code de la carte illustrant le workflow 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 illustrant le mappage de cette opération composite sur un élément fusionné dans LiteRT en utilisant la fonction en tant qu'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());
  }