Panoramica
Questa pagina descrive la progettazione e i passaggi necessari per convertire le operazioni composite in TensorFlow in operazioni fuse in TensorFlow Lite. Questa infrastruttura è di uso generale e supporta la conversione di qualsiasi operazione composita in TensorFlow in un'operazione fusa corrispondente in TensorFlow Lite.
Un esempio di utilizzo di questa infrastruttura è la fusione dell'operazione TensorFlow RNN in TensorFlow Lite, come dettagliato qui .
Cosa sono le operazioni fuse
Le operazioni di TensorFlow possono essere operazioni primitive, ad esempio tf.add , oppure possono essere composte da altre operazioni primitive, ad esempio tf.einsum . Un'operazione primitiva viene visualizzata come un singolo nodo nel grafico TensorFlow, mentre un'operazione composita è una raccolta di nodi nel grafico TensorFlow. L'esecuzione di un'operazione composita equivale all'esecuzione di ciascuna delle sue operazioni primitive costituenti.
Un'operazione fusa corrisponde a una singola operazione che sussume tutto il calcolo eseguito da ciascuna operazione primitiva all'interno dell'operazione composita corrispondente.
Vantaggi delle operazioni fuse
Esistono operazioni fuse per massimizzare le prestazioni delle implementazioni del kernel sottostanti, ottimizzando il calcolo generale e riducendo l'impronta di memoria. Questo è molto prezioso, soprattutto per carichi di lavoro di inferenza a bassa latenza e piattaforme mobili con risorse limitate.
Le operazioni fuse forniscono anche un'interfaccia di livello superiore per definire trasformazioni complesse come la quantizzazione, che altrimenti sarebbero impossibili o molto difficili da eseguire a un livello più granulare.
TensorFlow Lite ha molti casi di operazioni fuse per i motivi sopra articolati. Queste operazioni fuse corrispondono in genere alle operazioni composite nel programma TensorFlow di origine. Esempi di operazioni composite in TensorFlow implementate come un'unica operazione fusa in TensorFlow Lite includono varie operazioni RNN come sequenza unidirezionale e bidirezionale LSTM, convoluzione (conv2d, bias add, relu), completamente connesso (matmul, bias add, relu) e altro . In TensorFlow Lite, la quantizzazione LSTM è attualmente implementata solo nelle operazioni LSTM fuse.
Sfide con operazioni fuse
La conversione di operazioni composite da TensorFlow in operazioni fuse in TensorFlow Lite è un problema difficile. Questo è perché:
Le operazioni composite sono rappresentate nel grafico TensorFlow come un insieme di operazioni primitive senza un confine ben definito. Può essere molto difficile identificare (ad esempio tramite pattern matching) il sottografico corrispondente a tale operazione composita.
Potrebbe esserci più di un'implementazione TensorFlow destinata a un'operazione TensorFlow Lite fusa. Ad esempio, ci sono molte implementazioni LSTM in TensorFlow (Keras, Babelfish/lingvo ecc.) e ognuna di queste è composta da diverse operazioni primitive, ma tutte possono comunque essere convertite nella stessa operazione LSTM fusa in TensorFlow Lite.
In quanto tale, la conversione delle operazioni fuse si è rivelata piuttosto impegnativa.
Conversione da operazione composita a un'operazione personalizzata TFLite (consigliata)
Avvolgere l'operazione composita in una tf.function
In molti casi, alcune parti del modello possono essere mappate a una singola operazione in TFLite. Questo può aiutare con le prestazioni durante la scrittura di un'implementazione ottimizzata per operazioni specifiche. Per poter creare un'operazione fusa in TFLite, identificare la parte del grafico che rappresenta un'operazione fusa e racchiuderla in un tf.function con attributo "experimental_implements" in un tf.function , che ha valore di attributo tfl_fusable_op con valore true . Se l'operazione personalizzata accetta attributi, passali come parte degli stessi "implementi_sperimentali".
Esempio,
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))
Nota che non è necessario impostare allow_custom_ops sul convertitore poiché l'attributo tfl_fusable_op implica già.
Implementa l'operazione personalizzata e registrati con TFLite Interpreter
Implementare l'operazione fusa come operazione personalizzata TFLite: vedere le istruzioni .
Si noti che il nome con cui registrare l'operazione dovrebbe essere simile al nome specificato nell'attributo name nella firma degli attrezzi.
Un esempio per l'op nell'esempio è
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, ®);
Conversione dal funzionamento composito a quello fuso (avanzato)
L'architettura generale per la conversione delle operazioni composite TensorFlow in operazioni fuse TensorFlow Lite è la seguente:
Avvolgere l'operazione composita in una tf.function
Nel codice sorgente del modello TensorFlow, identifica e estrai l'operazione composita in una tf.function con l'annotazione della funzione experiment_implements . Vedere un esempio di ricerca incorporamento . La funzione definisce l'interfaccia e i suoi argomenti dovrebbero essere usati per implementare la logica di conversione.
Scrivi il codice di conversione
Il codice di conversione viene scritto per l'interfaccia della funzione con l'annotazione degli implements . Vedere un esempio di fusione per incorporare la ricerca . Concettualmente, il codice di conversione sostituisce l'implementazione composita di questa interfaccia con quella fusa.
Nel passaggio per le funzioni di preparazione del composito, inserisci il plug-in nel codice di conversione .
In usi più avanzati, è possibile implementare trasformazioni complesse degli operandi dell'operazione composita per derivare gli operandi dell'operazione fusa. Vedere Keras LSTM . codice di conversione come esempio.
Converti in TensorFlow Lite
Utilizza l'API TFLiteConverter.from_saved_model per convertire in TensorFlow Lite.
Sotto il cappuccio
Descriviamo ora dettagli di alto livello del progetto complessivo nella conversione in operazioni fuse in TensorFlow Lite.
Operazioni di composizione in TensorFlow
L'uso di tf.function con l'attributo della funzione Experiment_implements consente agli utenti di comporre in modo esplicito nuove operazioni utilizzando le operazioni primitive di TensorFlow e di specificare l'interfaccia implementata dall'operazione composita risultante. Questo è molto utile in quanto fornisce:
- Un confine ben definito per l'operazione composita nel grafico TensorFlow sottostante.
- Specificare in modo esplicito l'interfaccia implementata da questa operazione. Gli argomenti della
tf.functioncorrispondono agli argomenti di questa interfaccia.
Ad esempio, consideriamo un'operazione composita definita per implementare la ricerca di incorporamento. Questo si associa a un'operazione fusa in 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
Facendo in modo che i modelli utilizzino operazioni composite tramite tf.function come illustrato sopra, diventa possibile costruire un'infrastruttura generale per identificare e convertire tali operazioni in operazioni TensorFlow Lite fuse.
Estensione del convertitore TensorFlow Lite
Il convertitore TensorFlow Lite che è stato rilasciato all'inizio di quest'anno supportava solo l'importazione di modelli TensorFlow come grafico con tutte le variabili sostituite con i valori costanti corrispondenti. Questo non funziona per la fusione delle operazioni poiché tali grafici hanno tutte le funzioni allineate in modo che le variabili possano essere trasformate in costanti.
Per sfruttare la tf.function con la funzione experimental_implements durante il processo di conversione, le funzioni devono essere conservate fino a una fase successiva del processo di conversione.
Pertanto, abbiamo implementato un nuovo flusso di lavoro di importazione e conversione dei modelli TensorFlow nel convertitore per supportare il caso d'uso della fusione dell'operazione composita. Nello specifico, le novità introdotte sono:
- Importazione di modelli salvati di TensorFlow in MLIR
- fondere operazioni composite
- analisi di mutabilità variabile
- blocca tutte le variabili di sola lettura
Questo ci consente di eseguire la fusione dell'operazione utilizzando le funzioni che rappresentano le operazioni composite prima dell'integrazione della funzione e del congelamento delle variabili.
Implementazione dell'operazione di fusione
Diamo un'occhiata al passaggio di fusione dell'operazione in modo più dettagliato. Questo passaggio esegue le seguenti operazioni:
- Passa in rassegna tutte le funzioni nel modulo MLIR.
- Se una funzione ha l'attributo tf._implements, in base al valore dell'attributo, chiama l'utilità di fusione dell'operazione appropriata.
- L'utilità di fusione dell'operazione opera sugli operandi e sugli attributi della funzione (che fungono da interfaccia per la conversione) e sostituisce il corpo della funzione con un corpo di funzione equivalente contenente l'operazione fusa.
- In molti casi, il corpo sostituito conterrà operazioni diverse dall'operazione fusa. Questi corrispondono ad alcune trasformazioni statiche sugli operandi della funzione per ottenere gli operandi dell'operazione fusa. Poiché questi calcoli possono essere tutti ripiegati in modo costante, non sarebbero presenti nel flatbuffer esportato dove esisterebbe solo l'operazione fusa.
Ecco uno snippet di codice dal pass che mostra il flusso di lavoro principale:
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 */
}
Ecco un frammento di codice che mostra la mappatura di questa operazione composita in un'operazione fusa in TensorFlow Lite sfruttando la funzione come interfaccia di conversione.
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());
}