Panoramica
Questa pagina descrive la progettazione e i passaggi necessari per convertire le operazioni composte in TensorFlow alle operazioni unificate in LiteRT. Questa infrastruttura per uso generico e supporta la conversione di qualsiasi operazione composita in TensorFlow a un'operazione fusa corrispondente in LiteRT.
Un esempio di utilizzo di questa infrastruttura è TensorFlow RNN Operations Fusion LiteRT, come descritto qui.
Cosa sono le operazioni unificate
Le operazioni TensorFlow possono essere operazioni primitive, ad esempio tf.add oppure possono essere composte da altre operazioni primitive, ad esempio tf.einsum. Un modello viene visualizzata come singolo nodo nel grafico TensorFlow, mentre è una raccolta di nodi nel grafico TensorFlow. L'esecuzione di un un'operazione composita equivale all'esecuzione di ciascuna delle sue funzioni primitive operazioni aziendali.
Un'operazione combinata corrisponde a una singola operazione che include tutte le calcolo eseguito da ogni operazione primitiva all'interno dei un'operazione composita.
Vantaggi delle operazioni unificate
Esistono operazioni unite per massimizzare le prestazioni del kernel sottostante implementazioni, ottimizzando il calcolo complessivo e riducendo la memoria o meno. Questo è molto prezioso, soprattutto per i carichi di lavoro di inferenza a bassa latenza e le piattaforme per dispositivi mobili con risorse limitate.
Le operazioni combinate forniscono inoltre un'interfaccia di livello superiore per definire trasformazioni come la quantizzazione, che altrimenti sarebbe inattuabile o difficili da eseguire a un livello più granulare.
LiteRT ha molte istanze di operazioni unificate, per il motivo illustrate sopra. Queste operazioni unite corrispondono in genere a caratteristiche nel programma TensorFlow di origine. Esempi di operazioni composte in TensorFlow implementati come operazione unificata in LiteRT includono varie operazioni RNN, come la sequenza unidirezionale e bidirezionale, LSTM, convoluzione (conv2d, bias add, relu), completamente connessa (matmul, bias add, relu) e altro ancora. In LiteRT, la quantizzazione LSTM attualmente è disponibile solo implementato nelle operazioni LSTM integrate.
Sfide con operazioni integrate
Conversione delle operazioni composte da TensorFlow in operazioni unificate in LiteRT è un problema difficile. I motivi sono i seguenti:
Le operazioni composita sono rappresentate nel grafo TensorFlow come un insieme le operazioni primitive senza un confine ben definito. Può essere molto difficile da identificare (ad esempio tramite la corrispondenza di pattern) il sottografico corrispondente a un'operazione composita.
Potrebbero esserci più implementazioni di TensorFlow che hanno come target un insieme Operazione LiteRT. Ad esempio, esistono molte implementazioni LSTM in TensorFlow (Keras, Babelfish/lingvo ecc.) e ognuno di questi è composto operazioni primitive diverse, ma potrebbero comunque essere convertite la stessa operazione LSTM integrata in LiteRT.
Di conseguenza, la conversione delle operazioni fuse si è rivelata piuttosto impegnativa.
Conversione da un'operazione composita a un'operazione personalizzata TFLite (consigliata)
Aggrega l'operazione composita in un tf.function
In molti casi, parte del modello può essere mappata a una singola operazione
TFLite. Ciò può migliorare le prestazioni quando scrivi un'implementazione ottimizzata
per operazioni specifiche. Per poter creare un'operazione unificata in TFLite,
identificare la parte del grafico che rappresenta un'operazione fusa e includerla
un comando tf.function con
"implementazioni_sperimentali" a un tf.function, che dispone di un attributo
valore tfl_fusable_op con valore true. Se l'operazione personalizzata richiede
quindi passarli come parte dello stesso valore "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 sul convertitore come
L'attributo tfl_fusable_op implica già questo fatto.
Implementa un'operazione personalizzata e registrati con TFLite Interpreter
Implementare l'operazione combinata come operazione TFLite personalizzata - vedi istruzioni.
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 dell'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 unificata (avanzata)
L'architettura complessiva per convertire le operazioni composte di TensorFlow in Di seguito sono riportate le operazioni incluse con LiteRT:
Aggrega l'operazione composita in un tf.function
Nel codice sorgente del modello TensorFlow, identifica e astrai l'elemento composito in un tf.function con experimental_implements di una funzione. Guarda un esempio di ricerca di incorporamento. La definisce l'interfaccia e i suoi argomenti dovrebbero essere utilizzati per implementare della logica di conversione.
Scrivi il codice di conversione
Il codice di conversione è scritto in base all'interfaccia della funzione con la
Annotazione implements. Guarda un esempio di fusione per l'incorporamento
lookup. Concettualmente, il codice di conversione sostituisce il codice composito
di questa interfaccia con quella unificata.
Nel passaggio prepara le funzioni composte, inserisci il plug-in nella Google Cloud.
Negli utilizzi più avanzati, è possibile implementare trasformazioni complesse di gli operandi dell'operazione composita in modo da ricavare gli operandi operativa. Vedi Keras LSTM. del codice di conversione.
Converti in LiteRT
Utilizza la TFLiteConverter.from_saved_model API per la conversione in LiteRT.
dietro le quinte
Ora descriviamo i dettagli di alto livello del design complessivo della conversione in operazioni in LiteRT.
Operazioni di composizione in TensorFlow
L'utilizzo di tf.function con experimental_implements consente agli utenti di comporre in modo esplicito nuove operazioni utilizzando le operazioni primitive TensorFlow e specificare l'interfaccia alla quale viene viene implementata un'operazione composita. È molto utile in quanto fornisce:
- Un confine ben definito per l'operazione composita nell'area sottostante Grafico TensorFlow.
- Specifica in modo esplicito l'interfaccia implementata da questa operazione. La argomenti del tf.function corrispondono agli argomenti di questa interfaccia.
Consideriamo, ad esempio, un'operazione composita definita per implementare della ricerca di incorporamento. Questo viene mappato a un'operazione congiunta in LiteRT.
@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
I modelli utilizzano operazioni composte tramite tf.function come illustrato sopra, diventa possibile creare un'infrastruttura identificare e convertire queste operazioni in operazioni LiteRT integrate.
Estensione del convertitore LiteRT
Il convertitore LiteRT rilasciato all'inizio di quest'anno supportava solo importando i modelli TensorFlow come grafico con tutte le variabili sostituite con i rispettivi i valori costanti corrispondenti. Questa operazione non funziona per Operations Fusion poiché questi grafici hanno tutte le funzioni incorporate, in modo che le variabili possano essere costanti.
Per poter sfruttare
tf.function con
experimental_implements durante il processo di conversione, le funzioni
devono essere conservati fino a una fase successiva
del processo di conversione.
Pertanto, abbiamo implementato un nuovo flusso di lavoro di importazione e conversione di TensorFlow, nel convertitore per supportare il caso d'uso di fusione di operazioni composte. Nello specifico, le nuove funzionalità aggiunte sono:
- Importazione di modelli salvati di TensorFlow in MLIR
- operazioni composte
- analisi della mutabilità delle variabili
- blocca tutte le variabili di sola lettura
Questo ci permette di eseguire la fusione delle operazioni utilizzando le funzioni che rappresentano operazioni composte prima di eseguire l'incorporamento della funzione e il blocco delle variabili.
Implementazione di Operations Fusion
Diamo un'occhiata all'operazione Fusion pass in modo più dettagliato. Questa tessera seguenti:
- Scopri tutte le funzioni nel modulo MLIR.
- Se una funzione ha l'attributo tf._implements, in base all'attributo chiama l'utilità di fusione delle operazioni appropriata.
- L'utilità di fusione delle operazioni opera sugli operandi e (che fungono da interfaccia per la conversione) e sostituisce il corpo della funzione con un corpo funzione equivalente contenente un'operazione fusa.
- In molti casi, il corpo sostituito conterrà operazioni diverse dal un'operazione fusa. Queste corrispondono ad alcune trasformazioni statiche sul agli operandi di una funzione fusa per ottenere gli operandi dell'operazione fusa. Dal momento che questi calcoli possono essere sempre eliminati, non verrebbero nel flatbuffer esportato, dove solo l'operazione fused esistono.
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 in LiteRT 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());
}