Panoramica
In questa pagina vengono descritti la progettazione e i passaggi necessari per convertire le operazioni composite di TensorFlow in operazioni incorporate di TensorFlow Lite. Questa infrastruttura è per uso generico 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 delle operazioni TensorFlow RNN con TensorFlow Lite, come descritto qui.
Cosa sono le operazioni "fuse"
Le operazioni TensorFlow possono essere operazioni primitive, ad es. tf.add, o essere composte da altre operazioni primitive, ad esempio tf.einsum. Un'operazione primitiva viene visualizzata come 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 include tutti i calcoli eseguiti da ogni operazione primitiva all'interno della corrispondente operazione composita.
Vantaggi delle operazioni combinate
Esistono operazioni che consentono di massimizzare le prestazioni delle implementazioni sottostanti del kernel, ottimizzando il calcolo complessivo e riducendo l'ingombro della memoria. Ciò è molto utile, in particolare per i carichi di lavoro di inferenza a bassa latenza e per le piattaforme mobile con risorse limitate.
Le operazioni fusibili offrono inoltre un'interfaccia di livello superiore per definire trasformazioni complesse come la quantizzazione, che altrimenti non sarebbero fattibili o molto difficili da eseguire a un livello più granulare.
TensorFlow Lite ha molte istanze di operazioni combinate per i motivi illustrati sopra. In genere, queste operazioni combinate corrispondono a operazioni composite nel programma TensorFlow di origine. Esempi di operazioni composte in TensorFlow implementate come singola operazione fusa in TensorFlow Lite includono varie operazioni RNN come sequenza unidirezionale e bidirezionale LSTM, convoluzione (conv2d, bias add, relu), completamente connesse (matmul, bias add, relu) e altre ancora. In TensorFlow Lite, la quantizzazione di LSTM attualmente è implementata solo nelle operazioni LSTM fuse.
Sfide con operazioni fusi
La conversione delle operazioni composite da TensorFlow in operazioni fuse in TensorFlow Lite è un problema difficile. I motivi sono i seguenti:
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 es. tramite la corrispondenza di pattern) il grafico secondario corrispondente a un'operazione composita di questo tipo.
Potrebbe esserci più di un'implementazione TensorFlow che ha come target un'operazione incorporata di TensorFlow Lite. Ad esempio, esistono 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.
Di conseguenza, la conversione di operazioni fusibili si è rivelata piuttosto impegnativa.
Conversione da un'operazione composita a un'operazione personalizzata TFLite (consigliata)
Esegui il wrapping dell'operazione composita in un oggetto tf.function
In molti casi, una parte del modello può essere mappata a una singola operazione in
TFLite. Questo può favorire le prestazioni durante la scrittura
di un'implementazione ottimizzata per operazioni specifiche. Per poter creare un'operazione fusa in TFLite,
identifica la parte del grafico che rappresenta un'operazione fusa e includela in
una tf.function con
l'attributo "experimental_implements"
in un tf.function, che ha il valore dell'attributo tfl_fusable_op con valore true. Se l'operazione personalizzata accetta attributi, trasferiscili come parte dello stesso "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 impostare allow_custom_ops nel convertitore perché l'attributo tfl_fusable_op lo implica già.
Implementa un'operazione personalizzata e registrati con TFLite Interpreter
Implementa l'operazione combinata come operazione TFLite Custom; consulta le instructions.
Tieni presente che il nome con cui registrare l'operazione deve essere simile al nome
specificato nell'attributo name nella firma dell'implementazione.
Un esempio di operazione è
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 da operazione composita a operazione incorporata (avanzata)
Di seguito è riportata l'architettura complessiva per convertire le operazioni composite di TensorFlow in operazioni combinate di TensorFlow Lite:
Esegui il wrapping dell'operazione composita in un oggetto tf.function
Nel codice sorgente del modello TensorFlow, identifica ed estrae l'operazione composita in una funzione tf.function con l'annotazione della funzione experimental_implements. Guarda un esempio di embedding lookup. 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 è scritto in base all'interfaccia della funzione con l'annotazione implements. Visualizza una fusione di esempio per embedding
lookup. Concettualmente, il codice di conversione sostituisce l'implementazione
composita di questa interfaccia con quella fusa.
Nel pass prepare-composite-functions, inserisci il plug-in nel codice di conversione.
Negli usi più avanzati, è possibile implementare trasformazioni complesse degli operandi dell'operazione composta per ricavare gli operandi dell'operazione unita. Vedi Keras LSTM. come esempio.
Converti a TensorFlow Lite
Utilizza l'API TFLiteConverter.from_saved_model per la conversione in TensorFlow Lite.
dietro le quinte
Ora descriviamo i dettagli generali del design complessivo della conversione in operazioni combinate in TensorFlow Lite.
Composizione delle operazioni in TensorFlow
L'utilizzo di tf.function con l'attributo funzione experimental_implements consente agli utenti di scrivere esplicitamente nuove operazioni utilizzando le operazioni primitive di TensorFlow e specificare l'interfaccia implementata dall'operazione composita risultante. Questa funzionalità è molto utile in quanto fornisce:
- Un confine ben definito per l'operazione composita nel grafico TensorFlow sottostante.
- Specifica in modo esplicito l'interfaccia implementata da questa operazione. Gli argomenti di tf.function corrispondono agli argomenti di questa interfaccia.
Ad esempio, prendiamo in considerazione un'operazione composita definita per implementare la ricerca di incorporamento. che corrisponde 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
Se i modelli utilizzano operazioni composite tramite tf.function, come sopra illustrato, diventa possibile creare un'infrastruttura generale per identificare e convertire tali operazioni in operazioni di TensorFlow Lite fuse.
Estendere il convertitore TensorFlow Lite
Il convertitore TensorFlow Lite rilasciato all'inizio di quest'anno supportava solo l'importazione dei modelli TensorFlow sotto forma di grafico, con tutte le variabili sostituite dai rispettivi valori della costante corrispondente. Questo non funziona per la fusione operativa poiché questi grafici hanno tutte le funzioni incorporate in modo che le variabili possano essere trasformate in costanti.
Per utilizzare tf.function con la funzionalità experimental_implements durante il processo di conversione, è necessario che queste funzioni vengano conservate fino alle fasi successive 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 di fusione delle operazioni composite. In particolare, le nuove funzionalità aggiunte sono:
- Importazione dei modelli salvati in MLIR di TensorFlow
- fondere operazioni composite
- analisi della mutabilità delle variabili
- bloccare tutte le variabili di sola lettura
Questo ci consente di eseguire la fusione operativa utilizzando le funzioni che rappresentano le operazioni composite prima di funzionare con l'incorporamento e il blocco variabile.
Implementazione della fusione operativa
Vediamo più in dettaglio il passaggio di fusione dell'operazione. Questa tessera svolge le seguenti operazioni:
- Panoramica di 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 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 della 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 costantemente ripiegati, non sarebbero presenti nel flatbuffer esportato dove esiste solo l'operazione fusa.
Ecco uno snippet di codice della tessera 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 uno snippet di codice che mostra la mappatura di questa operazione composita a un'operazione combinata in TensorFlow Lite utilizzando 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());
}