Aperçu
Cette page décrit la conception et les étapes nécessaires pour convertir des 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 de l'opération TensorFlow RNN à TensorFlow Lite, comme détaillé ici .
Quelles sont les opérations fusionnées
Les opérations TensorFlow peuvent être soit des opérations primitives, par exemple tf.add , soit composées d'autres opérations primitives, par exemple tf.einsum . Une opération primitive apparaît sous la forme d'un nœud unique dans le graphe TensorFlow, tandis qu'une opération composite est une collection de nœuds dans le graphe TensorFlow. L'exécution d'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 subsume tous les calculs effectués par chaque opération primitive dans 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 à 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 à effectuer à un niveau plus granulaire.
TensorFlow Lite a de nombreuses instances d'opérations fusionnées pour les raisons énoncées ci-dessus. Ces opérations fusionnées correspondent généralement à des opérations composites dans le programme source TensorFlow. Des exemples d'opérations composites dans TensorFlow qui sont implémentées en une seule opération fusionnée dans TensorFlow Lite incluent diverses opérations RNN telles que la séquence unidirectionnelle et bidirectionnelle LSTM, la convolution (conv2d, bias add, relu), entièrement connectée (matmul, bias add, relu) et plus . 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:
Les opérations composites sont représentées dans le graphique TensorFlow sous la forme d'un ensemble d'opérations primitives sans limite bien définie. Il peut être très difficile d'identifier (par exemple via le pattern matching) le sous-graphe correspondant à une telle opération composite.
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 peuvent 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.
Conversion d'une opération composite en une opération personnalisée TFLite (recommandé)
Enveloppez l'opération composite dans un 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 encapsulez-la dans un tf.function avec l'attribut "experimental_implements" à 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émenter une opération personnalisée et s'inscrire 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 doit être similaire au nom spécifié dans l'attribut name dans la signature de l'outil.
Un exemple pour l'op 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 coder.
return kTfLiteOk;
};
reg.builtin_code = kTfLiteCustom;
resolver->AddCustom(kOpName, ®);
Conversion d'un fonctionnement composite en un fonctionnement fusionné (avancé)
L'architecture globale de conversion des opérations composites TensorFlow en opérations fusionnées TensorFlow Lite est la suivante :
Enveloppez l'opération composite dans un tf.function
Dans le code source du modèle TensorFlow, identifiez et extrayez l'opération composite dans une fonction tf.function avec l'annotation de fonction experimental_implements . Voir un exemple d' incorporation de lookup . 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 incorporer lookup . Conceptuellement, le code de conversion remplace l'implémentation composite de cette interface par celle fusionnée.
Dans la passe prepare-composite-functions, branchez 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 effectuer la conversion vers TensorFlow Lite.
Sous la capuche
Nous décrivons maintenant les détails de haut niveau de la conception globale lors 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 que l'opération composite résultante implémente. Ceci est très utile car il fournit:
- Une limite bien définie pour l'opération composite dans le graphique TensorFlow sous-jacent.
- Spécifiez explicitement l'interface que cette opération implémente. Les arguments de la
tf.functioncorrespondent 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 qui a été publié plus tôt cette année ne prenait en charge que l'importation de modèles TensorFlow sous forme de graphique, toutes les variables étant remplacées par leurs valeurs constantes correspondantes. Cela ne fonctionne pas pour la fusion d'opérations car de tels graphes ont toutes les fonctions alignées afin que les variables puissent être transformées en constantes.
Afin de tirer parti de la 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 place 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. Plus précisément, les nouvelles fonctionnalités ajoutées sont :
- Importation de modèles enregistrés TensorFlow dans MLIR
- fusionner les opérations composites
- analyse de mutabilité variable
- geler toutes les variables en lecture seule
Cela nous permet d'effectuer la fusion d'opérations en utilisant les fonctions représentant les opérations composites avant l'inlining des fonctions et le gel des variables.
Mise en œuvre de l'opération fusion
Examinons plus en détail la passe de fusion de l'opération. Ce laissez-passer effectue les opérations suivantes :
- Parcourez toutes les fonctions du module MLIR.
- Si une fonction possède l'attribut tf._implements, en fonction de la valeur de l'attribut, appelle l'utilitaire de fusion d'opération approprié.
- L'utilitaire de fusion d'opération 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.
- 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 constamment repliés, ils ne seraient pas présents dans le flatbuffer 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 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());
}