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 con 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. Eseguire un'operazione composita equivale a eseguire 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 della corrispondente operazione composita.
Vantaggi delle operazioni fuse
Le operazioni fuse esistono per massimizzare le prestazioni delle implementazioni del kernel sottostanti, ottimizzando il calcolo complessivo e riducendo l'ingombro della memoria. Ciò è molto utile, soprattutto per i carichi di lavoro di inferenza a bassa latenza e le piattaforme mobili con risorse limitate.
Le operazioni fuse forniscono anche un'interfaccia di livello superiore per definire trasformazioni complesse come la quantizzazione, che altrimenti sarebbe irrealizzabile o molto difficile da eseguire a un livello più granulare.
TensorFlow Lite presenta molte istanze 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 una singola operazione fusa in TensorFlow Lite includono varie operazioni RNN come sequenza unidirezionale e bidirezionale LSTM, convoluzione (conv2d, bias add, relu), completamente connessa (matmul, bias add, relu) e altro ancora . In TensorFlow Lite, la quantizzazione LSTM è attualmente implementata solo nelle operazioni LSTM fuse.
Sfide con operazioni fuse
La conversione delle operazioni composite da TensorFlow alle 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 patternmatching) il sottografo corrispondente a tale operazione composita.
Potrebbero esserci più implementazioni di TensorFlow destinate 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 potrebbero comunque essere convertite nella stessa operazione LSTM fusa in TensorFlow Lite.
Pertanto, la conversione delle operazioni fuse si è rivelata piuttosto impegnativa.
Conversione da operazione composita a operazione personalizzata TFLite (consigliata)
Racchiude l'operazione composita in una tf.function
In molti casi, alcune parti del modello possono essere mappate su una singola operazione in TFLite. Ciò può aiutare con le prestazioni quando si scrive 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 una tf.function con l'attributo "experimental_implements" in una tf.function , che ha valore di attributo tfl_fusable_op con valore true . Se l'operazione personalizzata accetta attributi, passali come parte degli stessi "experimental_implements".
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))
Tieni presente che non è necessario allow_custom_ops sul convertitore poiché l'attributo tfl_fusable_op lo implica già.
Implementa operazioni personalizzate e registrati con TFLite Interpreter
Implementa l'operazione fusa come operazione personalizzata TFLite: consulta le istruzioni .
Tieni presente che il nome con cui registrare l'operazione dovrebbe essere simile al nome specificato nell'attributo name nella firma dell'implementazione.
Un esempio per l'operazione 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 code.
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 di TensorFlow in operazioni fuse di TensorFlow Lite è riportata di seguito:
Racchiude l'operazione composita in una tf.function
Nel codice sorgente del modello TensorFlow, identifica e astratti l'operazione composita in una tf.function con l'annotazione della funzione sperimentale_implements . Guarda un esempio di ricerca di incorporamento . La funzione definisce l'interfaccia e i suoi argomenti devono essere utilizzati per implementare la logica di conversione.
Scrivi il codice di conversione
Il codice di conversione viene scritto tramite l'interfaccia della funzione con l'annotazione implements . Guarda un esempio di fusione per l'incorporamento della ricerca . Concettualmente, il codice di conversione sostituisce l'implementazione composita di questa interfaccia con quella fusa.
Nel passaggio prepare-composite-functions, inserisci il codice di conversione .
Negli usi più avanzati, è possibile implementare trasformazioni complesse degli operandi dell'operazione composta per derivare gli operandi dell'operazione fusa. Vedi 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 i dettagli di alto livello della progettazione complessiva nella conversione in operazioni fuse in TensorFlow Lite.
Composizione di operazioni in TensorFlow
L'uso di tf.function con l'attributo della funzione sperimentale_implements consente agli utenti di comporre in modo esplicito nuove operazioni utilizzando le operazioni primitive di TensorFlow e specificare l'interfaccia che implementa l'operazione composita risultante. Questo è molto utile in quanto fornisce:
- Un confine ben definito per l'operazione composita nel grafico TensorFlow sottostante.
- Specificare esplicitamente l'interfaccia implementata da questa operazione. Gli argomenti della
tf.functiontf.corrispondono 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 rilasciato all'inizio di quest'anno supportava solo l'importazione di modelli TensorFlow come grafico con tutte le variabili sostituite con i corrispondenti valori costanti. Questo non funziona per le operazioni di fusione poiché tali grafici hanno tutte le funzioni in linea in modo che le variabili possano essere trasformate in costanti.
Per sfruttare tf.function con la funzionalità experimental_implements durante il processo di conversione, le funzioni devono essere preservate fino a una fase successiva del processo di conversione.
Pertanto, abbiamo implementato un nuovo flusso di lavoro per l'importazione e la conversione dei modelli TensorFlow nel convertitore per supportare il caso d'uso della fusione operativa composita. Nello specifico le nuove funzionalità aggiunte sono:
- Importazione dei modelli salvati di TensorFlow in MLIR
- fondere le operazioni composite
- analisi di mutabilità variabile
- congelare tutte le variabili di sola lettura
Ciò ci consente di eseguire la fusione delle operazioni utilizzando le funzioni che rappresentano le operazioni composite prima dell'inlining delle funzioni e del congelamento delle variabili.
Implementazione dell'operazione fusione
Diamo un'occhiata più in dettaglio all'operazione Fusion Pass. Questo passaggio effettua le seguenti operazioni:
- Passa in rassegna tutte le funzioni del 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 delle operazioni opera sugli operandi e sugli attributi della funzione (che fungono da interfaccia per la conversione) e sostituisce il corpo della funzione con un corpo della funzione equivalente contenente l'operazione fusa.
- In molti casi, il corpo sostituito conterrà operazioni diverse da quella con il fuso. 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 lo snippet di codice del 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 lo snippet di codice che mostra la mappatura di questa operazione composita su 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());
}