Présentation
Cette page décrit la conception et les étapes requises pour convertir les opérations composites de TensorFlow en opérations fusionnées dans TensorFlow Lite. Cette infrastructure est à usage général et permet la conversion de toute opération composite dans TensorFlow en opération fusionnée correspondante dans TensorFlow Lite.
Un exemple d'utilisation de cette infrastructure est la fusion d'opérations de RNN TensorFlow avec TensorFlow Lite, comme indiqué ici.
Que sont les Fused Operations ?
Les opérations TensorFlow peuvent être des opérations primitives comme tf.add ou elles peuvent être composées d'autres opérations primitives comme 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 un ensemble de nœuds dans le graphe TensorFlow. L'exécution d'une opération composite équivaut à l'exécution de chacune de ses opérations primitives.
Une opération fusionnée correspond à une seule opération qui regroupe tous les calculs effectués par chaque opération primitive dans l'opération composite correspondante.
Avantages des opérations Fused
Les opérations fusionnées existent pour maximiser les performances des implémentations de noyau sous-jacentes, en optimisant le calcul global et en réduisant l'encombrement de la mémoire. Cela s'avère très utile, en particulier pour les charges de travail d'inférence à faible latence et les plates-formes mobiles limitées en ressources.
Les opérations fusionnées fournissent également une interface de niveau supérieur permettant de définir des transformations complexes telles que la quantification, qui seraient autrement impossibles à réaliser ou difficiles à réaliser à un niveau plus précis.
TensorFlow Lite comporte de nombreuses instances d'opérations fusionnées pour les raisons décrites ci-dessus. Ces opérations fusionnées correspondent généralement à des opérations composites du programme TensorFlow source. Des exemples d'opérations composites dans TensorFlow qui sont mises en œuvre en tant qu'opération fusionnée unique dans TensorFlow Lite incluent diverses opérations de RNN comme les séquences unidirectionnelles et bidirectionnelles LSTM, la convolution (conv2d, préjugé ajouté, relu), les connexions entièrement connectées (matmul, polarisation, relu), etc. Dans TensorFlow Lite, la quantification LSTM n'est actuellement mise en œuvre que dans les opérations Fused LSTM.
Défis liés aux Fused Operations
La conversion d'opérations composites de TensorFlow en opérations fusionnées dans TensorFlow Lite est un problème difficile. Voici pourquoi :
Les opérations composites sont représentées dans le graphe TensorFlow comme un ensemble d'opérations primitives sans limite bien définie. Il peut être très difficile d'identifier (par exemple, via la correspondance de modèle) 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.). Chacune d'elles est composée d'opérations primitives différentes, mais elles peuvent toutes être converties en la même opération LSTM fusionnée dans TensorFlow Lite.
La conversion d'opérations fusionnées s'est donc révélée assez difficile.
Convertir une opération composite en opération personnalisée TFLite (recommandé)
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, identifiez la partie du graphe qui représente une opération fusionnée, puis encapsulez-la dans une fonction tf.function avec l'attribut "experimental_Implements" pour tf.function, dont la valeur d'attribut tfl_fusable_op a la valeur true. Si l'opération personnalisée prend des attributs, puis les transmet dans le cadre du même élément "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, car l'attribut tfl_fusable_op l'implique déjà.
Implémenter une opération personnalisée et s'enregistrer avec l'interpréteur TFLite
Implémentez votre opération fusionnée en tant qu'opération personnalisée TFLite. Consultez les instructions.
Notez que le nom avec lequel enregistrer l'opération doit être semblable au nom spécifié dans l'attribut name de la signature.
L'opération présentée 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, ®);
Convertir une opération composite en opération fusionnée (avancé)
L'architecture globale de conversion des opérations composites TensorFlow en opérations fusionnées TensorFlow Lite est présentée ci-dessous:
Encapsuler 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. Consultez un exemple de recherche de représentations vectorielles continues. La fonction définit l'interface, et ses arguments doivent être utilisés pour implémenter la logique de conversion.
Rédiger le code de conversion
Le code de conversion est écrit conformément à l'interface de la fonction avec l'annotation implements. Consultez un exemple de fusion pour la recherche de représentations vectorielles continues. Sur le plan conceptuel, le code de conversion remplace l'implémentation composite de cette interface par l'implémentation fusionnée.
Lors de la transmission des fonctions composites, ajoutez un plug-in à votre code de conversion.
Dans des cas d'utilisation plus avancés, il est possible de mettre en œuvre des transformations complexes des opérandes de l'opération composite afin de dériver les opérandes de l'opération fusionnée. Consultez le code de conversion Keras LSTM à titre d'exemple.
Convertir au format TensorFlow Lite
Utilisez l'API TFLiteConverter.from_saved_model pour effectuer la conversion au format TensorFlow Lite.
dans le détail
Nous décrivons maintenant en détail 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 mise en œuvre par l'opération composite résultante. Cette fonctionnalité est très utile, car elle offre les avantages suivants:
- Limite bien définie pour l'opération composite dans le graphe TensorFlow sous-jacent.
- Spécifiez explicitement l'interface mise en œuvre par cette opération. Les arguments de tf.function correspondent à ceux de cette interface.
Prenons l'exemple d'une opération composite définie pour implémenter la recherche de représentations vectorielles continues. Cette valeur 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.
Étendre le convertisseur TensorFlow Lite
Le convertisseur TensorFlow Lite, lancé plus tôt cette année, n'acceptait que l'importation de modèles TensorFlow sous forme de graphe, dans laquelle toutes les variables étaient remplacées par les valeurs constantes correspondantes. Cela ne fonctionne pas pour la fusion des opérations, car ces graphiques ont toutes les fonctions intégrées afin que les variables puissent être transformées en constantes.
Pour pouvoir exploiter tf.function avec la fonctionnalité experimental_implements pendant le processus de conversion, les fonctions doivent être conservées jusqu'à une étape ultérieure du processus de conversion.
Ainsi, nous avons mis en œuvre un nouveau workflow d'importation et de conversion de modèles TensorFlow dans le convertisseur pour prendre en charge le cas d'utilisation de fusion d'opérations composites. Les nouvelles fonctionnalités ajoutées sont les suivantes:
- Importation de modèles enregistrés TensorFlow dans MLIR
- fusibles opérations composites
- analyse de mutabilité variable
- figer toutes les variables en lecture seule
Cela nous permet d'effectuer la fusion des opérations à l'aide des fonctions représentant les opérations composites avant l'intégration des fonctions et le gel des variables.
Implémenter la fusion d'opérations
Examinons plus en détail l'opération de fusion des opérations. Cette carte effectue les opérations suivantes:
- Effectuer une boucle dans 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érations approprié.
- L'utilitaire de fusion des opérations agit 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 é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 à certaines 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 pliés de manière constante, ils ne seront pas présents dans le flatbuffer exporté, où seule l'opération fusionnée existerait.
Voici un extrait de code de la carte montrant 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 avec 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());
}