Übersicht
Auf dieser Seite werden das Design und die Schritte beschrieben, die zum Konvertieren von zusammengesetzten Operationen in TensorFlow in zusammengeführte Operationen in LiteRT erforderlich sind. Diese Infrastruktur ist universell und unterstützt die Konvertierung beliebiger zusammengesetzter Operationen in TensorFlow in eine entsprechende zusammengeführte Operation in LiteRT.
Ein Beispiel für die Verwendung dieser Infrastruktur ist die TensorFlow-RNN-Vorgangsfusion zu LiteRT, wie hier beschrieben.
Was sind zusammengeführte Vorgänge?
TensorFlow-Operationen können entweder primitive Operationen sein, z.B. tf.add, oder aus anderen primitiven Operationen zusammengesetzt werden, z.B. tf.einsum. Ein einfacher Vorgang wird als einzelner Knoten im TensorFlow-Diagramm dargestellt, während ein zusammengesetzter Vorgang eine Sammlung von Knoten im TensorFlow-Diagramm ist. Die Ausführung einer zusammengesetzten Operation entspricht der Ausführung jeder ihrer primitiven Operationen.
Eine fusionierte Operation entspricht einer einzelnen Operation, die alle Berechnungen umfasst, die von jeder primitiven Operation innerhalb der entsprechenden zusammengesetzten Operation ausgeführt werden.
Vorteile von zusammengeführten Vorgängen
Zusammengeführte Vorgänge sind dazu da, die Leistung der zugrunde liegenden Kernel-Implementierungen zu maximieren, indem die Gesamtberechnung optimiert und der Speicherbedarf reduziert wird. Das ist besonders für Inferenz-Arbeitslasten mit niedriger Latenz und ressourcenbeschränkte mobile Plattformen sehr nützlich.
Zusammengeführte Vorgänge bieten auch eine übergeordnete Schnittstelle zum Definieren komplexer Transformationen wie der Quantisierung, die auf einer detaillierteren Ebene ansonsten nicht möglich oder sehr schwierig wären.
LiteRT enthält viele Instanzen von fusionierten Operationen aus den oben genannten Gründen. Diese zusammengeführten Vorgänge entsprechen in der Regel zusammengesetzten Vorgängen im TensorFlow-Quellprogramm. Beispiele für zusammengesetzte Vorgänge in TensorFlow, die in LiteRT als einzelner zusammengeführter Vorgang implementiert werden, sind verschiedene RNN-Vorgänge wie unidirektionale und bidirektionale Sequenz-LSTMs, Faltungen (conv2d, bias add, relu) und vollständig verbundene Vorgänge (matmul, bias add, relu). In LiteRT ist die LSTM-Quantisierung derzeit nur in den zusammengeführten LSTM-Operationen implementiert.
Herausforderungen bei zusammengeführten Vorgängen
Die Konvertierung zusammengesetzter Operationen von TensorFlow in zusammengeführte Operationen in LiteRT ist schwierig. Dies hat folgende Gründe:
Zusammengesetzte Vorgänge werden im TensorFlow-Diagramm als eine Reihe von primitiven Vorgängen ohne klar definierte Grenze dargestellt. Es kann sehr schwierig sein, den Untergraphen zu identifizieren, der einem solchen zusammengesetzten Vorgang entspricht (z.B. über Mustervergleich).
Es kann mehr als eine TensorFlow-Implementierung für einen zusammengeführten LiteRT-Vorgang geben. Es gibt beispielsweise viele LSTM-Implementierungen in TensorFlow (Keras, Babelfish/lingvo usw.). Jede davon besteht aus verschiedenen primitiven Operationen, aber alle könnten in LiteRT in denselben zusammengeführten LSTM-Vorgang konvertiert werden.
Die Umwandlung von zusammengeführten Vorgängen hat sich daher als sehr schwierig erwiesen.
Von einem zusammengesetzten Vorgang zu einem benutzerdefinierten TFLite-Vorgang konvertieren (empfohlen)
Schließen Sie den zusammengesetzten Vorgang in tf.function ein.
In vielen Fällen kann ein Teil des Modells einer einzelnen Operation in TFLite zugeordnet werden. Dies kann bei der Optimierung der Leistung helfen, wenn eine optimierte Implementierung für bestimmte Vorgänge geschrieben wird. Um einen zusammengeführten Vorgang in TFLite zu erstellen, identifizieren Sie den Teil des Diagramms, der einen zusammengeführten Vorgang darstellt, und umschließen Sie ihn mit einer tf.function mit dem Attribut „experimental_implements“ für ein tf.function, das den Attributwert tfl_fusable_op mit dem Wert true hat. Wenn für den benutzerdefinierten Vorgang Attribute erforderlich sind, übergeben Sie sie als Teil derselben „experimental_implements“.
Beispiel:
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))
Beachten Sie, dass Sie allow_custom_ops nicht für den Converter festlegen müssen, da dies bereits durch das Attribut tfl_fusable_op impliziert wird.
Benutzerdefinierten Vorgang implementieren und beim TFLite-Interpreter registrieren
Implementieren Sie den zusammengeführten Vorgang als benutzerdefinierten TFLite-Vorgang. Hier finden Sie eine Anleitung.
Der Name, mit dem der Vorgang registriert wird, sollte dem Namen ähneln, der im Attribut name in der Implementierungssignatur angegeben ist.
Ein Beispiel für den Vorgang im Beispiel ist
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, ®);
Von einem zusammengesetzten zu einem zusammengeführten Vorgang wechseln (Advanced)
Die allgemeine Architektur für die Konvertierung von zusammengesetzten TensorFlow-Operationen in zusammengeführte LiteRT-Operationen ist unten dargestellt:
Schließen Sie den zusammengesetzten Vorgang in tf.function ein.
Suchen Sie im TensorFlow-Modellquellcode nach dem zusammengesetzten Vorgang und abstrahieren Sie ihn in eine tf.function mit der Funktionsanmerkung experimental_implements. Beispiel für die Einbettungssuche Die Funktion definiert die Schnittstelle und ihre Argumente sollten verwendet werden, um die Conversion-Logik zu implementieren.
Conversion-Code schreiben
Der Conversion-Code wird gemäß der Schnittstelle der Funktion mit der Annotation implements geschrieben. Beispiel für die Fusion für die Einbettungssuche Konzeptionell ersetzt der Conversion-Code die zusammengesetzte Implementierung dieser Schnittstelle durch die zusammengeführte.
Fügen Sie im Durchlauf „prepare-composite-functions“ Ihren Conversion-Code ein.
Bei komplexeren Anwendungen ist es möglich, komplexe Transformationen der Operanden der zusammengesetzten Operation zu implementieren, um die Operanden der fusionierten Operation abzuleiten. Ein Beispiel für den Konvertierungscode finden Sie unter Keras-LSTM.
In LiteRT konvertieren
Verwenden Sie die API TFLiteConverter.from_saved_model, um in LiteRT zu konvertieren.
Funktionsweise
Im Folgenden werden die allgemeinen Details des Gesamtdesigns bei der Konvertierung in zusammengeführte Vorgänge in LiteRT beschrieben.
Vorgänge in TensorFlow zusammenstellen
Durch die Verwendung von tf.function mit dem Funktionsattribut experimental_implements können Nutzer explizit neue Vorgänge aus primitiven TensorFlow-Vorgängen zusammensetzen und die Schnittstelle angeben, die der resultierende zusammengesetzte Vorgang implementiert. Das ist sehr nützlich, da es Folgendes bietet:
- Eine klar definierte Grenze für den zusammengesetzten Vorgang im zugrunde liegenden TensorFlow-Graphen.
- Geben Sie explizit die Schnittstelle an, die von diesem Vorgang implementiert wird. Die Argumente von tf.function entsprechen den Argumenten dieser Schnittstelle.
Betrachten wir als Beispiel einen zusammengesetzten Vorgang, der für die Implementierung der Einbettungssuche definiert ist. Dies entspricht einer zusammengeführten Operation 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
Wenn Modelle wie oben gezeigt zusammengesetzte Vorgänge über tf.function verwenden, kann eine allgemeine Infrastruktur erstellt werden, um solche Vorgänge zu identifizieren und in zusammengeführte LiteRT-Vorgänge zu konvertieren.
LiteRT-Konverter erweitern
Der LiteRT-Konverter, der Anfang des Jahres veröffentlicht wurde, unterstützte nur den Import von TensorFlow-Modellen als Diagramm, in dem alle Variablen durch die entsprechenden konstanten Werte ersetzt wurden. Das funktioniert nicht für die Zusammenführung von Vorgängen, da in solchen Graphen alle Funktionen inline sind, sodass die Variablen in Konstanten umgewandelt werden können.
Damit die tf.function mit dem Feature experimental_implements während der Konvertierung verwendet werden kann, müssen die Funktionen bis später im Konvertierungsprozess beibehalten werden.
Daher haben wir einen neuen Workflow zum Importieren und Konvertieren von TensorFlow-Modellen im Konverter implementiert, um den Anwendungsfall für die Zusammenführung von zusammengesetzten Operationen zu unterstützen. Folgende Funktionen wurden hinzugefügt:
- TensorFlow-Modelle im SavedModel-Format in MLIR importieren
- Zusammengesetzte Vorgänge zusammenführen
- Analyse der Variabilität von Variablen
- Alle schreibgeschützten Variablen einfrieren
So können wir die Zusammenführung von Vorgängen mithilfe der Funktionen durchführen, die die zusammengesetzten Vorgänge vor dem Inlining von Funktionen und dem Einfrieren von Variablen darstellen.
Zusammenführung von Vorgängen implementieren
Sehen wir uns den Pass für das Zusammenführen von Vorgängen genauer an. Diese Karte bzw. dieses Ticket bietet Folgendes:
- Alle Funktionen im MLIR-Modul durchlaufen.
- Wenn eine Funktion das Attribut „tf._implements“ hat, wird basierend auf dem Attributwert das entsprechende Dienstprogramm für die Zusammenführung von Vorgängen aufgerufen.
- Das Tool zur Zusammenführung von Vorgängen arbeitet mit den Operanden und Attributen der Funktion (die als Schnittstelle für die Konvertierung dienen) und ersetzt den Funktionskörper durch einen gleichwertigen Funktionskörper, der den zusammengeführten Vorgang enthält.
- In vielen Fällen enthält der ersetzte Textkörper andere Vorgänge als den zusammengeführten Vorgang. Diese entsprechen einigen statischen Transformationen der Operanden der Funktion, um die Operanden des zusammengeführten Vorgangs zu erhalten. Da diese Berechnungen alle konstant gefaltet werden können, sind sie nicht im exportierten Flatbuffer vorhanden, in dem nur der zusammengeführte Vorgang vorhanden ist.
Hier ist ein Code-Snippet aus dem Ticket, das den Hauptworkflow zeigt:
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 */
}
Hier ist ein Code-Snippet, das zeigt, wie dieser zusammengesetzte Vorgang in LiteRT einem zusammengeführten Vorgang zugeordnet wird, wobei die Funktion als Konvertierungsschnittstelle dient.
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());
}