TensorFlow-Vorgangsfusion

Übersicht

Auf dieser Seite werden das Design und die Schritte beschrieben, die zum Umwandeln von zusammengesetzten Operationen erforderlich sind. in TensorFlow, um Operationen in LiteRT zu verschmelzen. Diese Infrastruktur ist für allgemeine Zwecke und unterstützt die Konvertierung beliebiger zusammengesetzter Operationen in TensorFlow. mit einer entsprechenden fusionierten Operation in LiteRT.

Ein Beispiel für diese Infrastruktur ist die Fusion von RNN-Operationen, LiteRT, wie hier beschrieben.

Was sind fusionierte Vorgänge

Zeichnen

TensorFlow-Vorgänge können entweder primitive Vorgänge sein, z.B. tf.add oder sie können aus anderen primitiven Operationen, z.B. tf.einsum an. Primitiv wird im TensorFlow-Diagramm als einzelner Knoten angezeigt, während eine zusammengesetzte ist eine Sammlung von Knoten in der TensorFlow-Grafik. Ausführen eines zusammengesetzte Operation entspricht der Ausführung jeder ihrer einzelnen Primitiven Geschäftsabläufe.

Eine fusionierte Operation entspricht einer einzelnen Operation, die alle die von jeder primitiven Operation in der entsprechenden zusammengesetzte Operation.

Vorteile fusionierter Vorgänge

Kombinierte Vorgänge dienen der Maximierung der Leistung des zugrunde liegenden Kernels Implementierungen durch Optimierung der gesamten Berechnung und Reduzierung des Arbeitsspeichers Fußabdruck. Dies ist sehr nützlich, insbesondere bei Inferenzarbeitslasten mit niedriger Latenz. und ressourcenbeschränkten mobilen Plattformen.

Kombinierte Operationen bieten auch eine höhere Schnittstelle zum Definieren komplexer wie Quantisierung, die ansonsten nicht umsetzbar wären oder auf detaillierterer Ebene schwer umzusetzen.

Bei LiteRT gibt es viele Instanzen von zusammengeführten Operationen. wie oben beschrieben. Diese zusammengeführten Vorgänge entsprechen in der Regel Operationen im TensorFlow-Quellprogramm. Beispiele für zusammengesetzte Operationen in TensorFlow, die als einzelner zusammengeführter Vorgang in LiteRT implementiert sind beinhalten verschiedene RNN-Operationen wie unidirektionale und bidirektionale Sequenzen LSTM, Faltung (Conv2d, Bias add, Relu), vollständig verbunden (Matmul, Bias add, und mehr. In LiteRT ist die LSTM-Quantisierung derzeit nur die in den fusionierten LSTM-Vorgängen implementiert sind.

Herausforderungen bei kombinierten Operationen

Zusammengesetzte Operationen von TensorFlow in fusionierte Operationen umwandeln in LiteRT ist ein schwieriges Problem. Dies hat folgende Gründe:

  1. Zusammengesetzte Operationen werden in der TensorFlow-Grafik als primitiven Operationen ohne eine genau definierte Grenze zu verstehen. Das kann sehr schwierig zu identifizieren (z.B. durch Musterabgleich) der Teilgrafik einer solchen zusammengesetzten Operation entspricht.

  2. Es kann mehr als eine TensorFlow-Implementierung für eine zusammengeführte LiteRT-Vorgang Es gibt viele LSTM-Implementierungen, in TensorFlow (Keras, Babelfish/lingvo usw.). Jedes dieser Elemente besteht aus primitive Operationen durchgeführt, aber alle konnten in die fusionierten LSTM-Vorgang in LiteRT.

Daher hat sich die Umwandlung fusionierter Vorgänge als schwierig erwiesen.

Zusammengesetzte Operation in tf.function zusammenfassen

In vielen Fällen kann ein Teil des Modells einem TFLite. Dies kann die Leistung verbessern, wenn Sie eine optimierte Implementierung schreiben für bestimmte Vorgänge. Um einen fusionierten Vorgang in TFLite erstellen zu können, identifizieren Sie den Teil des Graphen, der eine fusionierte Operation darstellt, eine tf.function mit „experimental_implements“ tf.function, das Attribut enthält, Wert tfl_fusable_op mit dem Wert true. Wenn der benutzerdefinierte Vorgang übergeben Sie sie 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 im Konverter nicht so festlegen: Das tfl_fusable_op-Attribut deutet das bereits an.

Benutzerdefinierte Operation implementieren und bei TFLite Interpreter registrieren

Implementieren Sie die zusammengeführte Operation als benutzerdefinierten TFLite-Vorgang: siehe Anleitung.

Der Name, mit dem der Vorgang registriert werden soll, sollte dem Namen ähneln. wird im Attribut name in der Implementierungssignatur angegeben.

Ein Beispiel für die Operation 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, &reg);

Zusammengesetzte in fusionierte Operation umwandeln (Advanced)

Die Gesamtarchitektur zum Konvertieren zusammengesetzter TensorFlow-Vorgänge in Kombinierte LiteRT-Operationen sind unten aufgeführt:

Zeichnen

Zusammengesetzte Operation in tf.function zusammenfassen

Identifizieren und abstrahieren Sie das zusammengesetzte Modell im Quellcode des TensorFlow-Modells. in eine tf.function mit dem Parameter experimental_implements Funktionsannotation. Ein Beispiel für eine Einbettungssuche ansehen. Die definiert die Schnittstelle und ihre Argumente sollten zur Implementierung der Conversion-Logik überprüfen.

Conversion-Code schreiben

Der Konvertierungscode wird über die Schnittstelle der Funktion mit dem implements-Anmerkung. Beispiel für eine Zusammenführung zum Einbetten ansehen Lookup Konzeptionell ersetzt der Conversion-Code das zusammengesetzte Implementierung dieser Schnittstelle mit der zusammengeführten Schnittstelle.

Im Pass „prep-composite-functions“ wird das Plug-in in der Conversion- Code hinzu.

Bei fortgeschritteneren Anwendungen ist es möglich, komplexe Transformationen von die Operanden der zusammengesetzten Operation, um die Operanden der zusammengeführten Operation abzuleiten. . Keras ansehen LSTM. Conversion-Code verwenden.

In LiteRT umwandeln

Verwenden Sie die Methode TFLiteConverter.from_saved_model API zum Konvertieren in LiteRT.

Details

Wir beschreiben nun allgemeine Details des Gesamtdesigns bei der Umwandlung in fusionierte Operationen in LiteRT an.

Vorgänge in TensorFlow erstellen

Die Verwendung von tf.function mit dem experimental_implements Funktionsattribut können Nutzende mithilfe von primitiven TensorFlow-Operationen ausführen und die Schnittstelle angeben, zusammengesetzte Operation implementiert. Dies ist sehr nützlich, da es Folgendes bietet:

  1. Eine klar definierte Grenze für die zusammengesetzte Operation in der zugrunde liegenden Operation TensorFlow-Grafik
  2. Geben Sie explizit die Schnittstelle an, die durch diesen Vorgang implementiert wird. Die die Argumente der tf.function entsprechen den Argumenten dieser Schnittstelle.

Sehen wir uns als Beispiel eine zusammengesetzte Operation an, die Embedding Lookup durch. 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

Indem Modelle zusammengesetzte Operationen über tf.function als wie oben dargestellt, wird es möglich, eine allgemeine Infrastruktur aufzubauen, Identifizieren und konvertieren Sie solche Vorgänge in zusammengeführte LiteRT-Vorgänge.

LiteRT-Konverter erweitern

Der LiteRT-Konverter, der Anfang des Jahres veröffentlicht wurde, Import von TensorFlow-Modellen als Graph, wobei alle Variablen durch ihre entsprechende konstante Werte. Dies funktioniert nicht für Operations Fusion, da Bei solchen Diagrammen sind alle Funktionen so angeordnet, dass die Variablen in Konstanten.

Um die Vorteile tf.function mit dem Parameter experimental_implements verwenden, werden die Funktionen müssen bis zu einem späteren Zeitpunkt im Konvertierungsprozess beibehalten werden.

Daher haben wir einen neuen Workflow für das Importieren und Konvertieren von TensorFlow im Converter so, dass sie den Anwendungsfall Zusammenführung von zusammengesetzten Operationen unterstützen. Folgende neue Funktionen wurden hinzugefügt:

  1. gespeicherte Modelle aus TensorFlow in MLIR importieren
  2. zusammengesetzte Vorgänge zusammenführen
  3. Variable Veränderlichkeitsanalyse
  4. alle schreibgeschützten Variablen fixieren

So können wir Operationen mit den Funktionen zusammenführen, zusammengesetzte Operationen vor der Inline-Funktion von Funktionen und dem Einfrieren von Variablen.

Operations Fusion implementieren

Sehen wir uns den Fusionspass des Vorgangs genauer an. Diese Karte bzw. dieses Ticket erfüllt Folgendes:

  1. Durchlaufen Sie alle Funktionen im MLIR-Modul.
  2. Wenn eine Funktion das Attribut „tf._implements“ basierend auf dem Attribut verwendet, wird das entsprechende Dienstprogramm für die Vorgangszusammenführung aufgerufen.
  3. Das Dienstprogramm „Operation Fusion“ arbeitet mit den Operanden der Funktion (die als Schnittstelle für die Conversion dienen) und ersetzt den Hauptteil der Funktion mit einem äquivalenten Funktionskörper, der die kombiniert.
  4. In vielen Fällen enthält der ersetzte Text andere Vorgänge als die kombiniert. Diese entsprechen einigen statischen Transformationen im der Operanden der Funktion, um die Operanden der fusionierten Operation zu erhalten. Da diese Berechnungen alle ständig zusammengefügt werden können, im exportierten Flatbuffer vorhanden, in dem nur der fusionierte Vorgang existieren.

Hier ist ein Code-Snippet aus der Karte, 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 sehen Sie ein Code-Snippet, das die Zuordnung dieses zusammengesetzten Vorgangs zu einer zusammengeführten Operation zeigt. in LiteRT und nutzt diese Funktion als Konvertierungsschnittstelle.

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());
  }