Überblick
Auf dieser Seite werden das Design und die Schritte beschrieben, die zum Konvertieren von zusammengesetzten Vorgängen in TensorFlow in fusionierte Vorgänge in TensorFlow Lite erforderlich sind. Diese Infrastruktur ist für allgemeine Zwecke gedacht und unterstützt die Konvertierung eines zusammengesetzten Vorgangs in TensorFlow in einen entsprechenden fusionierten Vorgang in TensorFlow Lite.
Ein Beispiel für die Verwendung dieser Infrastruktur ist die TensorFlow RNN-Vorgangsfusion mit TensorFlow Lite, wie hier beschrieben.
Was sind fusionierte Operationen?
TensorFlow-Vorgänge können entweder primitive Vorgänge wie tf.add sein oder aus anderen primitiven Vorgängen wie tf.einsum zusammengesetzt sein. Ein einfacher Vorgang wird in der TensorFlow-Grafik als einzelner Knoten angezeigt, während ein zusammengesetzter Vorgang eine Sammlung von Knoten in der TensorFlow-Grafik ist. Die Ausführung eines zusammengesetzten Vorgangs entspricht der Ausführung aller zugehörigen primitiven Vorgänge.
Eine fusionierte Operation entspricht einem einzelnen Vorgang, bei dem alle Berechnungen, die von jedem primitiven Vorgang ausgeführt werden, dem entsprechenden zusammengesetzten Vorgang zugeordnet werden.
Vorteile von Fused Operations
Es gibt fusionierte Vorgänge, um die Leistung der zugrunde liegenden Kernelimplementierungen zu maximieren, indem die Gesamtberechnung optimiert und der Arbeitsspeicherbedarf reduziert wird. Dies ist sehr nützlich, insbesondere für Inferenzarbeitslasten mit niedriger Latenz und ressourcenbeschränkte mobile Plattformen.
Zusammengeführte Operationen bieten auch eine Schnittstelle auf höherer Ebene, um komplexe Transformationen wie Quantisierung zu definieren, was auf detaillierter Ebene sonst nicht möglich oder nur sehr schwer auszuführen wäre.
In TensorFlow Lite gibt es aus den oben genannten Gründen viele Instanzen von fusionierten Vorgängen. Diese fusionierten Vorgänge entsprechen in der Regel zusammengesetzten Vorgängen im TensorFlow-Quellprogramm. Beispiele für zusammengesetzte Operationen in TensorFlow, die als einzelne fusionierte Operation in TensorFlow Lite implementiert werden, umfassen verschiedene RNN-Operationen wie unidirektionale und bidirektionale Sequenz LSTM, Faltung (conv2d, Bidirektionale Sequenz, Relu), vollständig verbunden (matmul, bias add, relu) und mehr. In TensorFlow Lite wird die LSTM-Quantisierung derzeit nur in den fusionierten LSTM-Vorgängen implementiert.
Herausforderungen bei Fused Operations
Das Konvertieren von zusammengesetzten Operationen von TensorFlow in fusionierte Operationen in TensorFlow Lite ist ein schwieriges Problem. Das hat folgende Gründe:
Zusammengesetzte Vorgänge werden in der TensorFlow-Grafik als eine Reihe primitiver Vorgänge ohne klar definierte Grenze dargestellt. Es kann sehr schwierig sein, die Teilgrafik zu identifizieren (z.B. durch Musterabgleich), die einem solchen zusammengesetzten Vorgang entspricht.
Es kann sein, dass mehrere TensorFlow-Implementierungen auf einen fusionierten TensorFlow Lite-Vorgang ausgerichtet sind. Es gibt beispielsweise viele LSTM-Implementierungen in TensorFlow (Keras, Babelfish/lingvo usw.), die jeweils aus verschiedenen primitiven Vorgängen bestehen, die jedoch alle in den gleichen fusionierten LSTM-Vorgang in TensorFlow Lite konvertiert werden können.
Daher hat sich die Umwandlung zusammengeführter Operationen als ziemlich schwierig erwiesen.
Von zusammengesetzter Operation in einen benutzerdefinierten TFLite-Vorgang konvertieren (empfohlen)
Zusammensetzen der zusammengesetzten Operation in tf.function
In vielen Fällen kann ein Teil des Modells einem einzelnen Vorgang in TFLite zugeordnet werden. Dies kann die Leistung beim Schreiben einer optimierten Implementierung für bestimmte Vorgänge verbessern. Damit Sie eine fusionierte Operation in TFLite erstellen können, müssen Sie den Teil des Diagramms identifizieren, der eine fusionierte Operation darstellt, und ihn in eine tf.function mit dem Attribut „experimental_implements“ in eine tf.function einbinden, die den Attributwert tfl_fusable_op mit dem Wert true hat. Wenn der benutzerdefinierte Vorgang Attribute annimmt, übergeben Sie diese als Teil desselben „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))
Sie müssen allow_custom_ops für den Converter nicht festlegen, da das tfl_fusable_op-Attribut bereits impliziert.
Benutzerdefinierte Vorgänge implementieren und mit TFLite Interpreter registrieren
Implementieren Sie den fusionierten Vorgang als benutzerdefinierten TFLite-Vorgang – siehe instructions.
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 in diesem 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 zusammengesetzten in einen fusionierten Vorgang konvertieren (erweitert)
Die Gesamtarchitektur zum Konvertieren von zusammengesetzten TensorFlow-Vorgängen in fusionierte TensorFlow Lite-Vorgänge sieht so aus:
Zusammensetzen der zusammengesetzten Operation in tf.function
Identifizieren Sie im Quellcode des TensorFlow-Modells den zusammengesetzten Vorgang und führen Sie ihn mit der Funktionsannotation experimental_implements in eine tf.function aus. Beispiel für einen eingebetteten Lookup Die Funktion definiert die Schnittstelle und ihre Argumente sollten verwendet werden, um die Konvertierungslogik zu implementieren.
Conversion-Code schreiben
Der Konvertierungscode wird über die Schnittstelle der Funktion mit der Annotation implements geschrieben. Fusionsbeispiel für die Einbettungssuche ansehen Konzeptionell ersetzt der Conversion-Code die zusammengesetzte Implementierung dieser Schnittstelle durch die kombinierte Implementierung.
Plug-in in der Vorbereitungs-Composite-Funktion in Ihrem Conversion-Code.
In fortgeschritteneren Anwendungen ist es möglich, komplexe Transformationen der Operanden der zusammengesetzten Operation zu implementieren, um die Operanden der kombinierten Operation abzuleiten. Sehen Sie sich als Beispiel den Conversion-Code Keras LSTM an.
In TensorFlow Lite konvertieren
Verwenden Sie die TFLiteConverter.from_saved_model API, um eine Konvertierung in TensorFlow Lite durchzuführen.
Details
Wir beschreiben jetzt eine detaillierte Beschreibung des gesamten Designs für die Umwandlung in zusammengeführte Vorgänge in TensorFlow Lite.
Zusammensetzungen in TensorFlow
Durch die Verwendung von tf.function mit dem Funktionsattribut experimental_implements können Nutzer mithilfe von einfachen TensorFlow-Operationen explizit neue Vorgänge erstellen und die Schnittstelle angeben, die von der resultierenden zusammengesetzten Operation implementiert wird. Dies ist sehr nützlich, da es Folgendes bietet:
- Eine genau definierte Grenze für den zusammengesetzten Vorgang in der zugrunde liegenden TensorFlow-Grafik.
- Geben Sie explizit die Schnittstelle an, die durch diesen Vorgang implementiert wird. Die Argumente der tf.function entsprechen den Argumenten dieser Schnittstelle.
Betrachten wir als Beispiel eine zusammengesetzte Operation, die zur Implementierung der Einbettungssuche definiert ist. Dies entspricht einem fusionierten Vorgang 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
Wenn Modelle zusammengesetzte Vorgänge über tf.function verwenden, wie oben dargestellt, ist es möglich, eine allgemeine Infrastruktur zu erstellen, um solche Vorgänge zu identifizieren und in fusionierte TensorFlow Lite-Vorgänge umzuwandeln.
TensorFlow Lite-Konverter erweitern
Der TensorFlow Lite-Konverter, der Anfang dieses Jahres veröffentlicht wurde, unterstützt nur den Import von TensorFlow-Modellen als Grafik, wobei alle Variablen durch die entsprechenden konstanten Werte ersetzt wurden. Dies funktioniert nicht für die Operation-Fusion, da in solchen Diagrammen alle Funktionen enthalten sind, sodass die Variablen in Konstanten umgewandelt werden können.
Damit tf.function mit der experimental_implements-Funktion während der Konvertierung genutzt werden kann, müssen die Funktionen bis zu einem späteren Zeitpunkt im Konvertierungsprozess beibehalten werden.
Daher haben wir einen neuen Workflow zum Importieren und Konvertieren von TensorFlow-Modellen im Converter implementiert, um den Anwendungsfall der Zusammenführung zusammengesetzter Vorgänge zu unterstützen. Das sind die neuen Funktionen:
- Gespeicherte TensorFlow-Modelle in MLIR importieren
- Zusammengesetzte Operationen zusammenführen
- Variable Veränderlichkeitsanalyse
- alle schreibgeschützten Variablen fixieren
Auf diese Weise können wir Vorgänge mit den Funktionen zusammenführen, die die zusammengesetzten Operationen darstellen, bevor die Funktionen inline eingefügt und die Variablen eingefroren werden.
Vorgangsfusion implementieren
Sehen wir uns die Operation Fusion Pass etwas genauer an. Diese Karte bzw. dieses Ticket erfüllt folgende Aufgaben:
- Alle Funktionen im MLIR-Modul durchlaufen.
- Wenn eine Funktion das Attribut tf._implements hat, wird basierend auf dem Attributwert das entsprechende Operation-Fusion-Dienstprogramm aufgerufen.
- Das Fusionsdienstprogramm arbeitet mit den Operanden und Attributen der Funktion (die als Schnittstelle für die Konvertierung dienen) und ersetzt den Hauptteil der Funktion durch einen entsprechenden Funktionsrumpf, der die fusionierte Operation enthält.
- In vielen Fällen enthält der ersetzte Text andere Vorgänge als die fusionierte Operation. Sie entsprechen einigen statischen Transformationen an den Operanden der Funktion, um die Operanden der zusammengeführten Operation zu erhalten. Da diese Berechnungen alle konstant weggeklappt sein können, wären sie im exportierten Flatbuffer nicht vorhanden, wo nur die fusionierte Operation existieren würde.
Hier ist das Code-Snippet aus der Karte bzw. 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 */
}
Das folgende Code-Snippet zeigt, wie diese zusammengesetzte Operation einem kombinierten Vorgang in TensorFlow Lite zugeordnet wird, der die Funktion als Konvertierungsschnittstelle nutzt.
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());
}