Einführung in den Cloud TPU v5e-Inferenzkonverter [Öffentliche Vorschau]

Einführung

Der Cloud TPU Inference Converter bereitet und optimiert ein TensorFlow 2-Modell (TF2) für TPU-Inferenz. Der Converter wird in einer lokalen oder TPU-VM-Shell ausgeführt. Die TPU-VM-Shell wird empfohlen, da darin die für den Konverter erforderlichen Befehlszeilentools vorinstalliert sind. Dafür wird ein exportiertes SavedModel ausgeführt und die folgenden Schritte ausgeführt:

TPU-Konvertierung: Dem Modell werden TPUPartitionedCall und andere TPU-Vorgänge hinzugefügt, damit es auf der TPU ausgeführt werden kann. Standardmäßig weist ein zur Inferenz exportiertes Modell keine derartigen Vorgänge auf und kann nicht auf der TPU bereitgestellt werden, selbst wenn es auf der TPU trainiert wurde.
Batchverarbeitung: Dem Modell werden Batchverarbeitungsvorgänge hinzugefügt, um die In-Graph-Batchverarbeitung für einen besseren Durchsatz zu ermöglichen.
BFloat16-Konvertierung: Es konvertiert das Datenformat des Modells von float32 in bfloat16, um die Rechenleistung zu verbessern und die Nutzung des High Bandwidth Memory (HBM) auf der TPU zu reduzieren.
E/A-Formoptimierung: Es optimiert die Tensorformen für die zwischen der CPU und TPU übertragenen Daten, um die Bandbreitennutzung zu verbessern.

Beim Exportieren eines Modells erstellen Nutzer Funktionsaliasse für alle Funktionen, die sie auf der TPU ausführen möchten. Sie übergeben diese Funktionen an den Converter und der Converter platziert sie auf der TPU und optimiert sie. Wenn diese Funktionen TPU-inkompatible Operationen enthalten, ist es möglich, die Soft-Device-Platzierung zu aktivieren, damit die inkompatiblen Operationen auf der CPU ausgeführt werden.

Der Cloud TPU Inference Converter ist als Docker-Image verfügbar, das in jeder Umgebung mit installiertem Docker ausgeführt werden kann.

Geschätzte Zeit für die Durchführung der oben genannten Schritte: ca. 20 bis 30 Minuten

Vorbereitung

Das Modell muss ein TF2-Modell sein und in das SavedModel-Format exportiert werden.
Das Modell muss einen Funktionsalias für die TPU-Funktion haben. Im Codebeispiel wird dies beschrieben. In den folgenden Beispielen wird tpu_func als TPU-Funktionsalias verwendet.
Achten Sie darauf, dass die CPU Ihres Computers AVX-Anweisungen (Advanced Vector eXtensions) unterstützt, da die Tensorflow-Bibliothek (die Abhängigkeit des Cloud TPU-Inferenzkonverters) für die Verwendung von AVX-Anweisungen kompiliert wurde. Die meisten CPUs haben die AVX-Unterstützung.
1. Sie können lscpu | grep avx ausführen, um zu prüfen, ob der AVX-Anweisungssatz unterstützt wird.

Hinweise

Führen Sie vor der Einrichtung die folgenden Schritte aus:

Neues Projekt erstellen: Wählen Sie in der Google Cloud Console auf der Seite für die Projektauswahl ein Cloud-Projekt aus oder erstellen Sie eines.
Cloud TPU-VM einrichten: Erstellen Sie eine neue Cloud TPU-VM mit der Google Cloud Console oder gcloud oder verwenden Sie eine vorhandene Cloud TPU-VM, um Inferenz mit dem konvertierten Modell auf der Cloud TPU-VM auszuführen.
- Achten Sie darauf, dass das Cloud TPU-VM-Image auf Tensorflow basiert. Beispiel: --version=tpu-vm-tf-2.11.0.
- Das konvertierte Modell wird auf diese Cloud TPU-VM geladen und bereitgestellt.
Prüfen Sie, ob Sie die Befehlszeilentools haben, die Sie für die Verwendung des Cloud TPU-Inferenzkonverters benötigen. Sie können das Google Cloud SDK und Docker lokal installieren oder eine Cloud TPU-VM verwenden, auf der diese Software standardmäßig installiert ist. Mithilfe dieser Tools können Sie mit dem Converter-Bild interagieren.

Mit dem folgenden Befehl können Sie eine SSH-Verbindung zur Cloud TPU-VM herstellen:
```
gcloud compute tpus tpu-vm ssh ${tpu-name} --zone ${zone} --project ${project-id}
```

Umgebung einrichten

Richten Sie Ihre Umgebung über Ihre TPU-VM-Shell oder Ihre lokale Shell ein.

TPU-VM-Shell

Führen Sie in der TPU-VM-Shell die folgenden Befehle aus, um die Nicht-Root-Docker-Nutzung zuzulassen:
```
sudo usermod -a -G docker ${USER}
newgrp docker
```
Initialisieren Sie die Helper für Docker-Anmeldedaten:
```
gcloud auth configure-docker \
  us-docker.pkg.dev
```

Lokale Shell

Richten Sie die Umgebung in Ihrer lokalen Shell mit den folgenden Schritten ein:

Installieren Sie das Cloud SDK, das das gcloud-Befehlszeilentool enthält.
Installieren Sie Docker:

Nutzung von Nicht-Root-Docker zulassen:

sudo usermod -a -G docker ${USER}
newgrp docker

Melden Sie sich in Ihrer Umgebung an:
```
gcloud auth login
```
Initialisieren Sie die Helper für Docker-Anmeldedaten:
```
gcloud auth configure-docker \
    us-docker.pkg.dev
```

Rufen Sie das Docker-Image des Inferenzkonverters ab:

  CONVERTER_IMAGE=us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
  docker pull ${CONVERTER_IMAGE}

Konverter-Bild

Das Bild dient zur einmaligen Modellkonvertierung. Legen Sie die Modellpfade fest und passen Sie die Konverter-Optionen an Ihre Anforderungen an. Im Abschnitt Nutzungsbeispiele werden mehrere gängige Anwendungsfälle beschrieben.

docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      function_alias: "tpu_func"
    }
    batch_options {
      num_batch_threads: 2
      max_batch_size: 8
      batch_timeout_micros: 5000
      allowed_batch_sizes: 2
      allowed_batch_sizes: 4
      allowed_batch_sizes: 8
      max_enqueued_batches: 10
    }
'

Inferenz mit dem konvertierten Modell in der Cloud TPU-VM

# Initialize the TPU
resolver = tf.distribute.cluster_resolver.TPUClusterResolver("local")
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

# Load the model
model = tf.saved_model.load(${CONVERTED_MODEL_PATH})

# Find the signature function for serving
serving_signature = 'serving_default' # Change the serving signature if needed
serving_fn = model.signatures[serving_signature]
# Run the inference using requests.
results = serving_fn(**inputs)
logging.info("Serving results: %s", str(results))

Anwendungsbeispiele

Funktionsalias für die TPU-Funktion hinzufügen

Suchen oder erstellen Sie eine Funktion in Ihrem Modell, die alles umfasst, was Sie auf der TPU ausführen möchten. Wenn @tf.function nicht vorhanden ist, fügen Sie sie hinzu.
Geben Sie beim Speichern des Modells SaveOptions wie unten an, um model.tpu_func einen Alias func_on_tpu zuzuweisen.
Sie können diesen Funktionsalias an den Converter übergeben.

class ToyModel(tf.keras.Model):
  @tf.function(
      input_signature=[tf.TensorSpec(shape=[None, 10], dtype=tf.float32)])
  def tpu_func(self, x):
    return x * 1.0

model = ToyModel()
save_options = tf.saved_model.SaveOptions(function_aliases={
    'func_on_tpu': model.tpu_func,
})
tf.saved_model.save(model, model_dir, options=save_options)

Modell mit mehreren TPU-Funktionen konvertieren

Sie können mehrere Funktionen auf der TPU platzieren. Erstellen Sie einfach mehrere Funktionsaliasse und übergeben Sie sie in converter_options_string an den Converter.

tpu_functions {
  function_alias: "tpu_func_1"
}
tpu_functions {
  function_alias: "tpu_func_2"
}

Quantisierung

Die Quantisierung ist ein Verfahren, mit dem die Genauigkeit der Zahlen reduziert wird, die zur Darstellung der Parameter eines Modells verwendet werden. Dies führt zu einer geringeren Modellgröße und einer schnelleren Berechnung. Ein quantisiertes Modell bietet Verbesserungen beim Inferenzdurchsatz sowie eine geringere Arbeitsspeichernutzung und Speichergröße, allerdings auf Kosten geringer Genauigkeitsrückgänge.

Die neue Quantisierungsfunktion nach dem Training, die auf TPU abzielt, basiert auf der ähnlichen Funktion in TensorFlow Lite, die für das Targeting von mobilen und Edge-Geräten verwendet wird. Weitere Informationen zur Quantisierung im Allgemeinen finden Sie im Dokument zu TensorFlow Lite.

Konzepte der Quantisierung

In diesem Abschnitt werden Konzepte beschrieben, die sich auf die Quantisierung mit dem Inferenzkonverter beziehen.

Konzepte im Zusammenhang mit anderen TPU-Konfigurationen (z. B. Segmente, Hosts, Chips und TensorCores) werden auf der Seite TPU-Systemarchitektur beschrieben.

Quantisierung nach dem Training (PTQ): PTQ ist ein Verfahren, mit dem die Größe und Rechenkomplexität eines neuronalen Netzwerkmodells reduziert werden, ohne seine Genauigkeit erheblich zu beeinträchtigen. PTQ wandelt die Gleitkommagewichtungen und -aktivierungen eines trainierten Modells in Ganzzahlen mit niedrigerer Genauigkeit um, z. B. 8-Bit- oder 16-Bit-Ganzzahlen. Dies kann zu einer erheblichen Reduzierung der Modellgröße und der Inferenzlatenz führen, während nur ein geringer Genauigkeitsverlust entsteht.
Kalibrierung: Der Kalibrierungsschritt zur Quantisierung ist der Prozess der Erfassung von Statistiken über den Wertebereich, den die Gewichtungen und Aktivierungen eines neuronalen Netzwerkmodells vornehmen. Anhand dieser Informationen werden die Quantisierungsparameter für das Modell bestimmt. Dies sind die Werte, die zum Konvertieren der Gleitkommagewichtungen und -Aktivierungen in Ganzzahlen verwendet werden.
Repräsentatives Dataset: Ein repräsentatives Dataset für die Quantisierung ist ein kleines Dataset, das die tatsächlichen Eingabedaten für das Modell darstellt. Es wird während des Kalibrierungsschritts der Quantisierung verwendet, um Statistiken zum Wertebereich zu erfassen, den die Gewichtungen und Aktivierungen des Modells benötigen. Das repräsentative Dataset sollte die folgenden Eigenschaften haben:
- Sie sollte die tatsächlichen Eingaben in das Modell während der Inferenz korrekt darstellen. Das bedeutet, dass sie den Wertebereich abdecken sollte, den das Modell wahrscheinlich in der realen Welt sieht.
- Sie sollte alle Zweige von Bedingungen (z. B. tf.cond) durchlaufen, falls vorhanden. Dies ist wichtig, da der Quantisierungsprozess alle möglichen Eingaben für das Modell verarbeiten kann, auch wenn sie nicht explizit im repräsentativen Dataset dargestellt werden.
- Sie sollte groß genug sein, um genügend Statistiken zu erfassen und Fehler zu reduzieren. Als Faustregel wird empfohlen, mehr als 200 repräsentative Beispiele zu verwenden.
Das repräsentative Dataset kann eine Teilmenge des Trainings-Datasets oder ein separates Dataset sein, das speziell dafür entwickelt wurde, die realen Eingaben in das Modell darzustellen. Welches Dataset verwendet werden soll, hängt von der jeweiligen Anwendung ab.
Static Range Quantization (SRQ): SRQ bestimmt den Wertebereich für die Gewichtungen und Aktivierungen eines neuronalen Netzwerkmodells einmal während des Kalibrierungsschritts. Das bedeutet, dass für alle Eingaben in das Modell derselbe Wertebereich verwendet wird. Dies kann weniger genau als eine dynamische Bereichsquantisierung ausfallen, insbesondere bei Modellen mit einem breiten Bereich von Eingabewerten. Die Quantisierung statischer Bereiche erfordert während der Laufzeit jedoch weniger Rechenvorgänge als die Quantisierung dynamischer Bereiche.

Dynamic Range Quantization (DRQ): DRQ bestimmt den Wertebereich für die Gewichtungen und Aktivierungen eines neuronalen Netzwerkmodells für jede Eingabe. Dadurch passt sich das Modell an den Wertebereich der Eingabedaten an, was die Genauigkeit verbessern kann. Allerdings erfordert die dynamische Bereichsquantisierung während der Laufzeit mehr Rechenvorgänge als die statische Bereichsquantisierung.

Feature	Quantisierung statischer Bereiche	Dynamische Bereichsquantisierung
Wertebereich	Wird bei der Kalibrierung einmal festgelegt	Für jede Eingabe festgelegt
Genauigkeit	Kann weniger genau sein, insbesondere bei Modellen mit einer Vielzahl von Eingabewerten	Kann genauer sein, insbesondere bei Modellen mit einer Vielzahl von Eingabewerten
Komplexität	Einfacher	Komplexer
Berechnung während der Laufzeit	Weniger Rechenleistung	Mehr Berechnungen

Nur gewichtete Quantisierung: Die reine Gewichtungsquantisierung ist eine Art von Quantisierung, bei der nur die Gewichtungen eines neuronalen Netzwerkmodells quantisiert werden, während die Aktivierungen im Gleitkommawert bleiben. Dies kann eine gute Option für Modelle sein, die empfindlich auf Genauigkeit reagieren, da dies dazu beitragen kann, die Genauigkeit des Modells aufrechtzuerhalten.

Quantisierung verwenden

Zum Anwenden der Quantisierung konfigurieren Sie QuantizationOptions und legen die Converter-Optionen fest. Wichtige Optionen:

Tags: Eine Sammlung von Tags, die den MetaGraphDef innerhalb des SavedModel identifizieren, der quantisiert werden soll. Du musst nicht angeben, ob du nur eine MetaGraphDef hast.
signatur_keys: Sequenz von Schlüsseln, die SignatureDef identifizieren, enthält Ein- und Ausgaben. Wenn nicht angegeben, wird ["serving_default"] verwendet.
quantization_method: Anzuwendende Quantisierungsmethode. Wenn nicht angegeben, wird die STATIC_RANGE-Quantisierung angewendet.
op_set: Sollte als XLA beibehalten werden. Dies ist derzeit die Standardoption. Eine Angabe ist nicht erforderlich.
repräsentativ_datasets: Geben Sie das Dataset an, das für die Kalibrierung der Quantisierungsparameter verwendet wird.

Repräsentatives Dataset erstellen

Ein repräsentatives Dataset ist im Wesentlichen eine Iteration von Stichproben. Ein Beispiel ist eine Map von: {input_key: input_value}. Beispiel:

representative_dataset = [{"x": tf.random.uniform(shape=(3, 3))}
                          for _ in range(256)]

Die repräsentativen Datasets sollten als TFRecord-Dateien mit der Klasse TfRecordRepresentativeDatasetSaver gespeichert werden, die derzeit im tf-nightly-pip-Paket verfügbar ist. Beispiel:

# Assumed tf-nightly installed.
import tensorflow as tf
representative_dataset = [{"x": tf.random.uniform(shape=(3, 3))}
                          for _ in range(256)]
tf.quantization.experimental.TfRecordRepresentativeDatasetSaver(
       path_map={'serving_default': '/tmp/representative_dataset_path'}
    ).save({'serving_default': representative_dataset})

Beispiele

Im folgenden Beispiel wird das Modell mit dem Signaturschlüssel serving_default und dem Funktionsalias tpu_func quantisiert:

docker run \
  --mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
  --mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
  ${CONVERTER_IMAGE} \
  --input_model_dir=/tmp/input \
  --output_model_dir=/tmp/output \
  --converter_options_string=' \
    tpu_functions { \
      function_alias: "tpu_func" \
    } \
    external_feature_configs { \
      quantization_options { \
        signature_keys: "serving_default" \
        representative_datasets: { \
          key: "serving_default" \
          value: { \
            tfrecord_file_path: "${TF_RECORD_FILE}" \
          } \
        } \
      } \
    } '

Soft-Device-Platzierung aktivieren

Durch die weiche Geräteplatzierung können TPU-inkompatible Vorgänge so verschoben werden, dass sie während der Laufzeit auf der CPU ausgeführt werden. Wenn das Modell eine TPU-inkompatible Operation innerhalb der TPU-Funktion hat, die nicht einfach getrennt werden kann, können Sie die Soft-Device-Platzierung aktivieren, indem Sie enable_soft_device_placement=True für die TPU-Funktion festlegen. Dieses Flag signalisiert der Laufzeit, nach TPU-inkompatiblen Vorgängen zu suchen und sie auf der CPU auszuführen.

Die weiche Geräteplatzierung ist auch für die Fehlerbehebung hilfreich. Damit können Sie die TPU-inkompatible tf.print in der TPU-Funktion verwenden, um die Werte der Tensoren auszugeben.

Die weiche Geräteplatzierung funktioniert am besten mit der MLIR-Bridge. Es muss ein zusätzliches Flag festgelegt werden.

tpu_functions {
  function_alias: "tpu_func"
  enable_soft_device_placement: true
}

Batchverarbeitung hinzufügen

Mit dem Converter können einem Modell Batchverarbeitung hinzugefügt werden. Eine Beschreibung der Batchoptionen, die abgestimmt werden können, finden Sie unter Definition von Batchoptionen.

Standardmäßig fasst der Converter alle TPU-Funktionen im Modell in Batches zusammen. Außerdem lassen sich damit von Nutzern bereitgestellte Signaturen und Funktionen in Batches zusammenfassen, wodurch die Leistung weiter verbessert werden kann. Jede TPU-Funktion, eine vom Nutzer bereitgestellte Funktion oder eine Signatur, die in Batches zusammengefasst ist, muss die strengen Formanforderungen des Batchvorgangs erfüllen.

Der Converter kann auch vorhandene Batchoptionen aktualisieren. Das folgende Beispiel zeigt, wie Sie einem Modell Batchverarbeitung hinzufügen. Weitere Informationen zur Batchverarbeitung finden Sie unter Batching im Detail.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}

bfloat16- und IO-Formoptimierung deaktivieren

BFloat16- und E/A-Formoptimierungen sind standardmäßig aktiviert. Wenn sie nicht gut mit Ihrem Modell funktionieren, können sie deaktiviert werden.

# Disable both optimizations
disable_default_optimizations: true

# Or disable them individually
io_shape_optimization: DISABLED
bfloat16_optimization: DISABLED

Conversion-Bericht

Sie finden diesen Conversion-Bericht im Protokoll, nachdem Sie den Inferenzkonverter ausgeführt haben. Ein Beispiel dafür sehen Sie unten.

-------- Conversion Report --------
TPU cost of the model: 96.67% (2034/2104)
CPU cost of the model:  3.33% (70/2104)

Cost breakdown
================================
%         Cost    Name
--------------------------------
3.33      70      [CPU cost]
48.34     1017    tpu_func_1
48.34     1017    tpu_func_2
--------------------------------

In diesem Bericht werden die Berechnungskosten des Ausgabemodells für CPU und TPU geschätzt. Außerdem werden die TPU-Kosten für jede Funktion aufgeschlüsselt. Dies sollte Ihre Auswahl der TPU-Funktionen in den Converter-Optionen widerspiegeln.

Wenn Sie die TPU besser nutzen möchten, sollten Sie mit der Modellstruktur experimentieren und die Converter-Optionen anpassen.

Häufig gestellte Fragen

Welche Funktion(en) soll ich auf der TPU platzieren?

Es empfiehlt sich, so viel von Ihrem Modell wie möglich auf der TPU zu speichern, da die überwiegende Mehrheit der Operationen auf der TPU schneller ausgeführt wird.

Wenn Ihr Modell keine TPU-inkompatiblen Operationen, Strings oder dünnbesetzten Tensoren enthält, ist normalerweise die beste Strategie das gesamte Modell auf der TPU zu platzieren. Dazu suchen oder erstellen Sie eine Funktion, die das gesamte Modell umschließt, erstellen einen Funktionsalias dafür und übergeben diesen an den Converter.

Wenn Ihr Modell Teile enthält, die nicht mit der TPU verwendet werden können (z.B. TPU-inkompatible Vorgänge, Strings oder dünnbesetzte Tensoren), hängt die Auswahl der TPU-Funktionen davon ab, wo sich der inkompatible Teil befindet.

Wenn es sich am Anfang oder am Ende des Modells befindet, können Sie das Modell refaktorieren, damit es auf der CPU bleibt. Beispiele sind Vor- und Nachbearbeitungsphasen von Strings. Weitere Informationen zum Verschieben von Code zur CPU finden Sie unter Wie verschiebe ich einen Teil des Modells zur CPU? Sie zeigt eine typische Refaktorierung des Modells.
Wenn es sich in der Mitte des Modells befindet und nur wenige Vorgänge sind, kann eine Soft-Device-Platzierung in Betracht gezogen werden. Aufgrund der zusätzlichen Kommunikation zwischen CPU und TPU kann es jedoch zu einer höheren Inferenzlatenz kommen.
Wenn es sich in der Mitte des Modells befindet und viele Vorgänge erforderlich sind, würde die weiche Geräteplatzierung die Inferenz zu kostspielig machen. Daher ist es besser, das Modell in drei Teile aufzuteilen und alle TPU-inkompatiblen Vorgänge im mittleren Teil aufzunehmen und auf der CPU auszuführen.
Wenn es ein dünnbesetzter Tensor ist, sollten Sie tf.sparse.to_dense für die CPU aufrufen und den resultierenden dichten Tensor an den TPU-Teil des Modells übergeben.

Ein weiterer zu berücksichtigender Faktor ist die HBM-Nutzung. Für das Einbetten von Tabellen kann viel HBM verwendet werden. Wenn sie die Hardwarebeschränkung der TPU überschreiten, müssen sie zusammen mit den Lookup-Operationen auf die CPU gesetzt werden.

Wann immer möglich sollte nur eine TPU-Funktion unter einer Signatur vorhanden sein. Wenn für die Struktur Ihres Modells mit jeder eingehenden Inferenzanfrage mehrere TPU-Funktionen aufgerufen werden müssen, sollten Sie sich der zusätzlichen Latenz bewusst sein, die durch das Senden von Tensoren zwischen CPU und TPU entsteht.

Der Conversion-Bericht bietet eine gute Möglichkeit, die Auswahl von TPU-Funktionen zu bewerten. Sie zeigt den Prozentsatz der auf der TPU platzierten Berechnungen und eine Aufschlüsselung der Kosten für jede TPU-Funktion.

Wie verschiebe ich einen Teil des Modells zur CPU?

Wenn Ihr Modell Teile enthält, die nicht auf der TPU bereitgestellt werden können, müssen Sie das Modell refaktorieren, um sie auf die CPU zu verschieben. Hier ist ein Beispiel für ein Spielzeug. Das Modell ist ein Sprachmodell mit einer Vorverarbeitungsphase. Der Code für Ebenendefinitionen und -funktionen wurde der Einfachheit halber weggelassen.

class LanguageModel(tf.keras.Model):
  @tf.function
  def model_func(self, input_string):
    word_ids = self.preprocess(input_string)
    return self.bert_layer(word_ids)

Dieses Modell kann aus zwei Gründen nicht direkt auf der TPU bereitgestellt werden. Der Parameter ist ein String. Zweitens kann die Funktion preprocess viele Stringvorgänge enthalten. Beide sind nicht TPU-kompatibel.

Zum Refaktorieren dieses Modells können Sie eine weitere Funktion namens tpu_func zum Hosten des rechenintensiven bert_layer erstellen. Erstellen Sie dann einen Funktionsalias für tpu_func und übergeben Sie ihn an den Converter. Auf diese Weise wird alles in tpu_func auf der TPU und alles, was in model_func übrig ist, auf der CPU ausgeführt.

class LanguageModel(tf.keras.Model):
  @tf.function
  def tpu_func(self, word_ids):
    return self.bert_layer(word_ids)

  @tf.function
  def model_func(self, input_string):
    word_ids = self.preprocess(input_string)
    return self.tpu_func(word_ids)

Was sollte ich tun, wenn das Modell TPU-inkompatible Operationen, Strings oder dünnbesetzte Tensoren hat?

Die meisten TensorFlow-Standardvorgänge werden auf der TPU unterstützt. Einige, einschließlich dünnbesetzter Tensoren und Strings, werden jedoch nicht unterstützt. Der Converter prüft nicht auf TPU-inkompatible Vorgänge. Ein Modell, das solche Vorgänge enthält, kann also die Konvertierung übergeben. Wenn es jedoch für die Inferenz ausgeführt wird, treten Fehler wie unten auf.

'tf.StringToNumber' op isn't compilable for TPU device. enable soft_device_placement option to run on CPU

Wenn Ihr Modell TPU-inkompatible Vorgänge enthält, sollten diese außerhalb der TPU-Funktion platziert werden. Darüber hinaus ist der String ein nicht unterstütztes Datenformat auf der TPU. Variablen vom Typ „String“ sollten daher nicht in der TPU-Funktion platziert werden. Außerdem sollten die Parameter und Rückgabewerte der TPU-Funktion nicht ebenfalls vom Typ String sein. Vermeiden Sie es in ähnlicher Weise, dünnbesetzte Tensoren in der TPU-Funktion zu platzieren, einschließlich ihrer Parameter und Rückgabewerte.

Es ist normalerweise nicht schwierig, den inkompatiblen Teil des Modells zu refaktorieren und auf die CPU zu verschieben. Hier ist ein Beispiel. Es gibt jedoch Fälle, in denen TPU-inkompatible Operationen tief im Modell verschachtelt und schwer zu trennen sind. Für diese Modelle können Sie die Soft-Device-Platzierung aktivieren, um die Vorgänge zur Laufzeit auf der CPU auszuführen. Beachten Sie jedoch die Auswirkungen auf die Leistung.

Wie unterstützt man benutzerdefinierte Vorgänge im Modell?

Wenn in Ihrem Modell benutzerdefinierte Vorgänge verwendet werden, erkennt der Converter sie möglicherweise nicht und konvertiert das Modell nicht. Dies liegt daran, dass die Vorgangsbibliothek des benutzerdefinierten Vorgangs, die die vollständige Definition des Vorgangs enthält, nicht mit dem Converter verknüpft ist.

Da der Converter-Code derzeit noch nicht als Open Source verfügbar ist, kann er nicht mit einem benutzerdefinierten Vorgang erstellt werden.

Was sollte ich tun, wenn ich ein TensorFlow 1-Modell habe?

Der Converter unterstützt keine TensorFlow 1-Modelle. TensorFlow 1-Modelle sollten zu TensorFlow 2 migriert werden.

Muss ich die MLIR-Bridge aktivieren, wenn ich mein Modell ausführe?

Die meisten konvertierten Modelle können entweder mit der neueren TF2XLA-MLIR-Bridge oder der ursprünglichen TF2XLA-Bridge ausgeführt werden.

Sowohl während des Trainings als auch der Inferenz funktioniert die Soft-Device-Platzierung in TensorFlow 2 am besten mit der MLIR-Brücke. Die MLIR-Bridge kann mit dieser Anleitung aktiviert werden.

Wie konvertiere ich ein Modell, das bereits ohne Funktionsalias exportiert wurde?

Wenn ein Modell ohne Funktionsalias exportiert wurde, ist es am einfachsten, es noch einmal zu exportieren und einen Funktionsalias zu erstellen. Wenn die Option nicht noch einmal exportiert werden kann, können Sie das Modell dennoch konvertieren, indem Sie eine concrete_function_name angeben. Allerdings erfordert die Identifizierung des richtigen concrete_function_name einige Detektivarbeit.

Funktionsaliasse sind eine Zuordnung von einem benutzerdefinierten String zu einem konkreten Funktionsnamen. Sie erleichtern es, auf eine bestimmte Funktion im Modell zu verweisen. Der Converter akzeptiert sowohl Funktionsaliasse als auch unbearbeitete konkrete Funktionsnamen.

Konkrete Funktionsnamen finden Sie unter saved_model.pb.

Das folgende Beispiel zeigt, wie eine konkrete Funktion namens __inference_serve_24 auf die TPU angewendet wird.

sudo docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      concrete_function_name: "__inference_serve_24"
    }'

Wie behebe ich einen Fehler bei der Beschränkung der Kompilierungszeitkonstante?

XLA erfordert sowohl für das Training als auch für die Inferenz, dass die Eingaben an bestimmte Operationen zum Zeitpunkt der TPU-Kompilierung eine bekannte Form haben. Wenn also XLA den TPU-Teil des Programms kompiliert, müssen die Eingaben an diese Operationen eine statisch bekannte Form haben.

Es gibt zwei Möglichkeiten, dieses Problem zu beheben.

Die beste Option besteht darin, die Eingaben des Vorgangs zum Zeitpunkt der Kompilierung des TPU-Programms durch XLA auf eine statisch bekannte Form zu aktualisieren. Diese Kompilierung erfolgt, bevor der TPU-Teil des Modells ausgeführt wird. Das bedeutet, dass die Form zum Zeitpunkt der Ausführung des TpuFunction statisch bekannt sein sollte.
Eine weitere Option besteht darin, TpuFunction so zu ändern, dass der problematische Vorgang nicht mehr einbezogen wird.

Warum wird ein Fehler bei der Batchverarbeitung für Formen angezeigt?

Für die Batchverarbeitung gelten strikte Formanforderungen, sodass eingehende Anfragen entlang ihrer 0. Dimension (auch als Batching-Dimension bezeichnet) gruppiert werden können. Diese Formanforderungen stammen aus dem TensorFlow-Batchvorgang und können nicht gelockert werden.

Andernfalls treten folgende Fehler auf:

Batch-Eingabetensoren müssen mindestens eine Dimension haben.
Die Dimensionen der Eingaben sollten übereinstimmen.
Die in einem bestimmten Vorgangsaufruf angegebenen Batch-Eingabetensoren müssen dieselbe Größe der 0. Dimension haben.
Die 0. Dimension des Batch-Ausgabetensors entspricht nicht der Summe der 0. Dimensionsgrößen der Eingabetensoren.

Ziehen Sie in Betracht, eine andere Funktion oder Signatur für einen Batch bereitzustellen, um diese Anforderungen zu erfüllen. Es kann auch erforderlich sein, vorhandene Funktionen zu ändern, um diese Anforderungen zu erfüllen.

Wenn eine Funktion in Batches zusammengefasst wird, muss die Form von „input_signature“ von @tf.function in der 0. Dimension „None“ enthalten. Wenn eine Signatur im Batch zusammengefasst wird, achten Sie darauf, dass alle Eingaben -1 in der 0. Dimension haben.

Eine vollständige Erklärung, warum diese Fehler auftreten und wie Sie sie beheben können, finden Sie unter Batch-Analyse.

Bekannte Probleme

TPU-Funktion kann nicht indirekt eine andere TPU-Funktion aufrufen

Der Converter kann die meisten Szenarien für Funktionsaufrufe über die CPU-TPU-Grenze hinweg verarbeiten. Es gibt jedoch einen seltenen Grenzfall, bei dem er fehlschlägt. Dabei wird von einer TPU-Funktion indirekt eine andere TPU-Funktion aufgerufen.

Dies liegt daran, dass der Converter den direkten Aufrufer einer TPU-Funktion so ändert, dass statt der TPU-Funktion selbst ein TPU-Aufruf-Stub aufgerufen wird. Der Aufruf-Stub enthält Operationen, die nur mit der CPU arbeiten können. Wenn eine TPU-Funktion eine Funktion aufruft, die schließlich den direkten Aufrufer aufruft, können diese CPU-Vorgänge auf der TPU ausgeführt werden, wodurch fehlende Kernelfehler generiert werden. Dieser Fall unterscheidet sich von einer TPU-Funktion, die direkt eine andere TPU-Funktion aufruft. In diesem Fall ändert der Converter keine der Funktionen so, dass der Aufruf-Stub aufgerufen wird, sodass es funktionieren kann.

Im Konverter haben wir die Erkennung dieses Szenarios implementiert. Wenn der folgende Fehler angezeigt wird, hat Ihr Modell diesen Grenzfall erreicht:

Unable to place both "__inference_tpu_func_2_46" and "__inference_tpu_func_4_68"
on the TPU because "__inference_tpu_func_2_46" indirectly calls
"__inference_tpu_func_4_68". This behavior is unsupported because it can cause
invalid graphs to be generated.

Die allgemeine Lösung besteht darin, das Modell zu refaktorieren, um ein Szenario für den Funktionsaufruf zu vermeiden. Wenn Sie dies schwierig finden, wenden Sie sich bitte an das Google-Supportteam, um mehr zu erfahren.

Referenz

Konvertierungsoptionen im Protobuf-Format

message ConverterOptions {
  // TPU conversion options.
  repeated TpuFunction tpu_functions = 1;

  // The state of an optimization.
  enum State {
    // When state is set to default, the optimization will perform its
    // default behavior. For some optimizations this is disabled and for others
    // it is enabled. To check a specific optimization, read the optimization's
    // description.
    DEFAULT = 0;
    // Enabled.
    ENABLED = 1;
    // Disabled.
    DISABLED = 2;
  }

  // Batch options to apply to the TPU Subgraph.
  //
  // At the moment, only one batch option is supported. This field will be
  // expanded to support batching on a per function and/or per signature basis.
  //
  //
  // If not specified, no batching will be done.
  repeated BatchOptions batch_options = 100;

  // Global flag to disable all optimizations that are enabled by default.
  // When enabled, all optimizations that run by default are disabled. If a
  // default optimization is explicitly enabled, this flag will have no affect
  // on that optimization.
  //
  // This flag defaults to false.
  bool disable_default_optimizations = 202;

  // If enabled, apply an optimization that reshapes the tensors going into
  // and out of the TPU. This reshape operation improves performance by reducing
  // the transfer time to and from the TPU.
  //
  // This optimization is incompatible with input_shape_opt which is disabled.
  // by default. If input_shape_opt is enabled, this option should be
  // disabled.
  //
  // This optimization defaults to enabled.
  State io_shape_optimization = 200;

  // If enabled, apply an optimization that updates float variables and float
  // ops on the TPU to bfloat16. This optimization improves performance and
  // throughtput by reducing HBM usage and taking advantage of TPU support for
  // bfloat16.
  //
  // This optimization may cause a loss of accuracy for some models. If an
  // unacceptable loss of accuracy is detected, disable this optimization.
  //
  // This optimization defaults to enabled.
  State bfloat16_optimization = 201;

  BFloat16OptimizationOptions bfloat16_optimization_options = 203;

  // The settings for XLA sharding. If set, XLA sharding is enabled.
  XlaShardingOptions xla_sharding_options = 204;
}

message TpuFunction {
  // The function(s) that should be placed on the TPU. Only provide a given
  // function once. Duplicates will result in errors. For example, if
  // you provide a specific function using function_alias do not also provide the
  // same function via concrete_function_name or jit_compile_functions.
  oneof name {
    // The name of the function alias associated with the function that
    // should be placed on the TPU. Function aliases are created during model
    // export using the tf.saved_model.SaveOptions.
    //
    // This is a recommended way to specify which function should be placed
    // on the TPU.
    string function_alias = 1;

    // The name of the concrete function that should be placed on the TPU. This
    // is the name of the function as it found in the GraphDef and the
    // FunctionDefLibrary.
    //
    // This is NOT the recommended way to specify which function should be
    // placed on the TPU because concrete function names change every time a
    // model is exported.
    string concrete_function_name = 3;

    // The name of the signature to be placed on the TPU. The user must make
    // sure there is no TPU-incompatible op under the entire signature. Or soft
    // device placement can be enabled. But this will incur performance loss.
    string signature_name = 5;

    // When jit_compile_functions is set to True, all jit compiled functions
    // are placed on the TPU.
    //
    // To use this option, decorate the relevant function(s) with
    // @tf.function(jit_compile=True), before exporting. Then set this flag to
    // True. The converter will find all functions that were tagged with
    // jit_compile=True and place them on the TPU.
    //
    // When using this option, all other settings for the TpuFunction such as
    // enable_soft_device_placement will apply to all functions tagged with
    // jit_compile=True.
    //
    // This option will place all jit_compile=True functions on the TPU.
    // If only some jit_compile=True functions should be placed on the TPU,
    // use function_alias or concrete_function_name.
    bool jit_compile_functions = 4;
  }

  // Set to true to enable outside compilation for this TPU function. If the TPU
  // function has TPU-incompatible ops, outside compilation can automatically
  // move the ops to execute on the CPU at the runtime.
  bool enable_soft_device_placement = 2;
}

message BatchOptions {
  // Number of scheduling threads for processing batches of work. Determines
  // the number of batches processed in parallel. This should be roughly in line
  // with the number of TPU cores available.
  int32 num_batch_threads = 1;

  // The maximum allowed batch size.
  int32 max_batch_size = 2;

  // Maximum number of microseconds to wait before outputting an incomplete
  // batch.
  int32 batch_timeout_micros = 3;

  // Optional list of allowed batch sizes. If left empty,
  // does nothing. Otherwise, supplies a list of batch sizes, causing the op
  // to pad batches up to one of those sizes. The entries must increase
  // monotonically, and the final entry must equal max_batch_size.
  repeated int32 allowed_batch_sizes = 4;

  // Maximum number of batches enqueued for processing before requests are
  // failed fast.
  int32 max_enqueued_batches = 5;

  // If set, disables large batch splitting which is an efficiency improvement
  // on batching to reduce padding inefficiency.
  bool disable_large_batch_splitting = 6;

  // Experimental features of batching. Everything inside is subject to change.
  message Experimental {
    // The component to be batched.
    // 1. Unset if it's for all TPU subgraphs.
    // 2. Set function_alias or concrete_function_name if it's for a function.
    // 3. Set signature_name if it's for a signature.
    oneof batch_component {
      // The function alias associated with the function. Function alias is
      // created during model export using the tf.saved_model.SaveOptions, and is
      // the recommended way to specify functions.
      string function_alias = 1;

      // The concreate name of the function. This is the name of the function as
      // it found in the GraphDef and the FunctionDefLibrary. This is NOT the
      // recommended way to specify functions, because concrete function names
      // change every time a model is exported.
      string concrete_function_name = 2;

      // The name of the signature.
      string signature_name = 3;
    }
  }

  Experimental experimental = 7;
}

message BFloat16OptimizationOptions {
  // Indicates where the BFloat16 optimization should be applied.
  enum Scope {
    // The scope currently defaults to TPU.
    DEFAULT = 0;
    // Apply the bfloat16 optimization to TPU computation.
    TPU = 1;
    // Apply the bfloat16 optimization to the entire model including CPU
    // computations.
    ALL = 2;
  }

  // This field indicates where the bfloat16 optimization should be applied.
  //
  // The scope defaults to TPU.
  Scope scope = 1;

  // If set, the normal safety checks are skipped. For example, if the model
  // already contains bfloat16 ops, the bfloat16 optimization will error because
  // pre-existing bfloat16 ops can cause issues with the optimization. By
  // setting this flag, the bfloat16 optimization will skip the check.
  //
  // This is an advanced feature and not recommended for almost all models.
  //
  // This flag is off by default.
  bool skip_safety_checks = 2;

  // Ops that should not be converted to bfloat16.
  // Inputs into these ops will be cast to float32, and outputs from these ops
  // will be cast back to bfloat16.
  repeated string filterlist = 3;
}

message XlaShardingOptions {
  // num_cores_per_replica for TPUReplicateMetadata.
  //
  // This is the number of cores you wish to split your model into using XLA
  // SPMD.
  int32 num_cores_per_replica = 1;

  // (optional) device_assignment for TPUReplicateMetadata.
  //
  // This is in a flattened [x, y, z, core] format (for
  // example, core 1 of the chip
  // located in 2,3,0 will be stored as [2,3,0,1]).
  //
  // If this is not specified, then the device assignments will utilize the same
  // topology as specified in the topology attribute.
  repeated int32 device_assignment = 2;

  // A serialized string of tensorflow.tpu.TopologyProto objects, used for
  // the topology attribute in TPUReplicateMetadata.
  //
  // You must specify the mesh_shape and device_coordinates attributes in
  // the topology object.
  //
  // This option is required for num_cores_per_replica > 1 cases due to
  // ambiguity of num_cores_per_replica, for example,
  // pf_1x2x1 with megacore and df_1x1
  // both have num_cores_per_replica = 2, but topology is (1,2,1,1) for pf and
  // (1,1,1,2) for df.
  // - For pf_1x2x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,2,1,1]
  //   device_coordinates=flatten([0,0,0,0], [0,1,0,0])
  // - For df_1x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,1,1,2]
  //   device_coordinates=flatten([0,0,0,0], [0,0,0,1])
  // - For df_2x2, mesh shape and device_coordinates looks like:
  //   mesh_shape = [2,2,1,2]
  //   device_coordinates=flatten(
  //    [0,0,0,0],[0,0,0,1],[0,1,0,0],[0,1,0,1]
  //    [1,0,0,0],[1,0,0,1],[1,1,0,0],[1,1,0,1])
  bytes topology = 3;
}

Batch-Erstellung im Detail

Die Batchverarbeitung wird verwendet, um den Durchsatz und die TPU-Auslastung zu verbessern. So können mehrere Anfragen gleichzeitig verarbeitet werden. Während des Trainings kann die Batchverarbeitung mit tf.data erfolgen. Während der Inferenz erfolgt dies in der Regel, indem eine Operation in die Grafik eingefügt wird, mit der eingehende Anfragen in Batches zusammengefasst werden. Der Vorgang wartet, bis er genügend Anfragen hat oder ein Zeitlimit erreicht ist, bevor er einen großen Batch aus den einzelnen Anfragen generiert. Weitere Informationen zu den verschiedenen Batchoptionen, die abgestimmt werden können, einschließlich Batchgrößen und Zeitüberschreitungen, finden Sie unter Definition der Batchoptionen.

Batchverarbeitung in Diagrammen

Standardmäßig fügt der Converter den Batchvorgang direkt vor der TPU-Berechnung ein. Sie umschließt die vom Nutzer bereitgestellten TPU-Funktionen und alle bereits vorhandenen TPU-Berechnungen im Modell in Batchverarbeitungsvorgängen. Sie können dieses Standardverhalten überschreiben, indem Sie dem Converter mitteilen, welche Funktionen und/oder Signaturen zusammengefasst werden sollen.

Das folgende Beispiel zeigt, wie die Standard-Batchverarbeitung hinzugefügt wird.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}

Batchverarbeitung von Signaturen

Bei der Batchverarbeitung für Signaturen wird das gesamte Modell von den Eingaben der Signatur bis zu den Ausgaben der Signatur in Batches zusammengefasst. Im Gegensatz zum Standard-Batchverarbeitungsverhalten des Converters werden bei der Signatur-Batchverarbeitung sowohl die TPU- als auch die CPU-Berechnung in Batches zusammengefasst. Dies führt bei einigen Modellen zu einer Leistungssteigerung von 10% bis 20% während der Inferenz.

Wie bei allen Batches gibt es auch für die Batchverarbeitung für Signaturen strikte Formanforderungen. Damit diese Formanforderungen erfüllt werden, sollten Signatureingaben Formen mit mindestens zwei Dimensionen haben. Die erste Dimension ist die Batchgröße und sollte eine Größe von -1 haben. Zum Beispiel sind (-1, 4), (-1) oder (-1, 128, 4, 10) gültige Eingabeformen. Wenn dies nicht möglich ist, können Sie die Standard-Batchverarbeitung oder die Funktions-Batchverarbeitung verwenden.

Wenn Sie Signatur-Batching verwenden möchten, geben Sie die Signaturnamen als signature_name mit BatchOptions an.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
  experimental {
    signature_name: "serving_default"
  }
}

Batchverarbeitung von Funktionen

Mit der Batchverarbeitung kann dem Converter mitgeteilt werden, welche Funktionen in Batches zusammengefasst werden sollen. Standardmäßig fasst der Converter alle TPU-Funktionen in Batches zusammen. Durch die Batchverarbeitung von Funktionen wird dieses Standardverhalten überschrieben.

Die Batchverarbeitung von Funktionen kann zur Batch-CPU-Berechnung verwendet werden. Bei vielen Modellen wird die Leistung durch die Batchverarbeitung der CPU-Berechnungen verbessert. Die beste Methode für die Batch-CPU-Berechnung ist die Verwendung der Signatur-Batchverarbeitung. Diese funktioniert jedoch möglicherweise bei einigen Modellen nicht. In diesen Fällen kann die Stapelverarbeitung von Funktionen verwendet werden, um zusätzlich zur TPU-Berechnung einen Teil der CPU-Berechnung zu verarbeiten. Der Batchverarbeitungsvorgang kann nicht auf der TPU ausgeführt werden. Daher muss jede bereitgestellte Batchfunktion auf der CPU aufgerufen werden.

Die Batchverarbeitung von Funktionen kann auch verwendet werden, um die strengen Formanforderungen zu erfüllen, die vom Batchvorgang auferlegt werden. Wenn die TPU-Funktionen die Formanforderungen des Batchverarbeitungsvorgangs nicht erfüllen, kann der Converter angewiesen werden, verschiedene Funktionen im Batch zu verarbeiten.

Generieren Sie dazu ein function_alias für die Funktion, die im Batch zusammengefasst werden soll. Suchen Sie dazu in Ihrem Modell nach einer Funktion oder erstellen Sie eine Funktion, die alles zusammenfasst, was Sie in Batches zusammenfassen möchten. Achten Sie darauf, dass diese Funktion die strengen Formanforderungen des Batchvorgangs erfüllt. Fügen Sie @tf.function hinzu, falls noch keine vorhanden sind. Es ist wichtig, die input_signature für @tf.function anzugeben. Die 0. Dimension sollte None sein, da es sich um die Batch-Dimension handelt und sie daher keine feste Größe haben kann. [None, 4], [None] oder [None, 128, 4, 10] sind beispielsweise gültige Eingabeformen. Geben Sie beim Speichern des Modells SaveOptions wie unten gezeigt an, um model.batch_func den Alias „batch_func“ zu verleihen. Anschließend können Sie diesen Funktionsalias an den Converter übergeben.

class ToyModel(tf.keras.Model):
  @tf.function(input_signature=[tf.TensorSpec(shape=[None, 10],
                                              dtype=tf.float32)])
  def batch_func(self, x):
    return x * 1.0

  ...

model = ToyModel()
save_options = tf.saved_model.SaveOptions(function_aliases={
    'batch_func': model.batch_func,
})
tf.saved_model.save(model, model_dir, options=save_options)

Übergeben Sie als Nächstes die function_alias mit den BatchOptions.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
  experimental {
    function_alias: "batch_func"
  }
}

Definition von Batchoptionen

num_batch_threads: (Ganzzahl) Anzahl der Planungsthreads für die Verarbeitung von Arbeitsstapeln. Bestimmt die Anzahl der parallel verarbeiteten Batches. Dies sollte ungefähr der Anzahl der verfügbaren TPU-Kerne entsprechen.
max_batch_size: (Ganzzahl) Maximal zulässige Batchgröße. Kann größer als allowed_batch_sizes sein, um eine große Batchaufteilung zu nutzen.
batch_timeout_micros: (Ganzzahl) Maximale Wartezeit in Mikrosekunden, bevor ein unvollständiger Batch ausgegeben wird.
allowed_batch_sizes: (Liste mit Ganzzahlen) Wenn die Liste nicht leer ist, werden Batches bis zur nächsten Größe in der Liste aufgefüllt. Die Liste muss kontinuierlich zunehmen und das letzte Element muss kleiner oder gleich max_batch_size sein.
max_enqueued_batches: (Ganzzahl) Maximale Anzahl von Batches, die zur Verarbeitung in die Warteschlange gestellt werden, bevor Anfragen schnell fehlschlagen.

Vorhandene Batching-Optionen aktualisieren

Sie können Batchoptionen hinzufügen oder aktualisieren, indem Sie das Docker-Image ausführen und dafür „batch_options“ angeben und disable_default_optimizations mit dem Flag --converter_options_string auf „true“ setzen. Die Batchoptionen werden auf jede TPU-Funktion oder jeden bereits vorhandenen Batchvorgang angewendet.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}
disable_default_optimizations=True

Formanforderungen im Batch

Batches werden durch Verkettung von Eingabetensoren über Anfragen hinweg entlang ihrer (0.) Batchdimension erstellt. Die Ausgabetensoren werden entlang ihrer 0. Dimension aufgeteilt. Für diese Vorgänge gelten strenge Formanforderungen für die Ein- und Ausgaben des Batchverarbeitungsvorgangs.

Schritt-für-Schritt-Anleitung

Zum Verständnis dieser Anforderungen ist es hilfreich zu verstehen, wie die Batchverarbeitung durchgeführt wird. Im folgenden Beispiel wird ein einfacher tf.matmul-Vorgang als Batch ausgeführt.

def my_func(A, B)
    return tf.matmul(A, B)

Die erste Inferenzanfrage erzeugt die Eingaben A und B mit den Formen (1, 3, 2) bzw. (1, 2, 4). Die zweite Inferenzanfrage erzeugt die Eingaben A und B mit den Formen (2, 3, 2) und (2, 2, 4).

Inferenzanfrage 1

Das Zeitlimit für die Batchverarbeitung wurde erreicht. Das Modell unterstützt eine Batchgröße von 3, sodass die Inferenzanfragen 1 und 2 ohne Auffüllung zusammengefasst werden. Die Batch-Tensoren werden durch Verkettung der Anfragen 1 und 2 entlang der Batchdimension (0.) gebildet. Da das A von 1 die Form (1, 3, 2) und das A von 2 die Form (2, 3, 2) hat, ist die resultierende Form (3, 3, 2), wenn sie entlang der Batchdimension (0.) verkettet werden.

Batchanfrage

Die tf.matmul wird ausgeführt und generiert eine Ausgabe mit der Form (3, 3, 4).

Batch-MatMul-Anfrage

Die Ausgabe von tf.matmul erfolgt als Batch, sodass sie wieder in separate Anfragen aufgeteilt werden muss. Der Batchvorgang führt dies durch eine Aufteilung entlang der Batchdimension (0.) jedes Ausgabetensors. Anhand der Form der ursprünglichen Eingaben wird entschieden, wie die 0. Dimension aufgeteilt wird. Da die Formen der Anfrage 1 die 0. Dimension von 1 haben, hat die Ausgabe bei der Form (1, 3, 4) eine 0. Dimension von 1. Da die Formen der Anfrage 2 die 0. Dimension von 2 haben, hat die Ausgabe eine 0. Dimension von 2 bei einer Form von (2, 3, 4).

Ergebnisse der Inferenzanfrage

Formanforderungen

Zum Ausführen der oben beschriebenen Eingabeverkettung und Ausgabeaufteilung gelten für den Batchvorgang die folgenden Formanforderungen:

Eingaben für die Batchverarbeitung dürfen keine Skalare sein. Für eine Verkettung entlang der 0. Dimension müssen die Tensoren mindestens zwei Dimensionen haben.

Gehen Sie in der obigen Schritt-für-Schritt-Anleitung vor. Weder A noch B sind Skalare.

Wenn diese Anforderung nicht erfüllt wird, führt dies zu einem Fehler wie: Batching input tensors must have at least one dimension. Eine einfache Lösung für diesen Fehler besteht darin, den Skalar zu einem Vektor zu machen.
Bei verschiedenen Inferenzanfragen (z. B. bei verschiedenen Sitzungsausführungsaufrufen) haben Eingabetensoren mit demselben Namen für jede Dimension mit Ausnahme der 0. Dimension dieselbe Größe. Auf diese Weise können Eingaben über ihre 0. Dimension verkettet werden.

In der obigen Schritt-für-Schritt-Anleitung hat das A der Anfrage 1 die Form (1, 3, 2). Das bedeutet, dass bei jeder zukünftigen Anfrage eine Form mit dem Muster (X, 3, 2) erzeugt werden muss. Anfrage 2 erfüllt diese Anforderung mit (2, 3, 2). Ähnlich hat die Form B der Anfrage 1 die Form (1, 2, 4), sodass alle zukünftigen Anfragen eine Form mit dem Muster (X, 2, 4) erzeugen müssen.

Wenn diese Anforderung nicht erfüllt wird, führt dies zu einem Fehler wie: Dimensions of inputs should match.
Bei einer bestimmten Inferenzanfrage müssen alle Eingaben dieselbe 0-te Dimensionsgröße haben. Wenn verschiedene Eingabetensoren für den Batchverarbeitungsvorgang unterschiedliche 0. Dimensionen haben, weiß der Batchverarbeitungsvorgang nicht, wie die Ausgabetensoren aufgeteilt werden sollen.

In der obigen Schritt-für-Schritt-Anleitung haben die Tensoren von Anfrage 1 alle die 0. Dimensionsgröße von 1. Dadurch weiß der Batchverarbeitungsvorgang, dass seine Ausgabe eine 0. Dimensionsgröße von 1 haben sollte. Entsprechend haben die Tensoren von Anfrage 2 die 0. Dimensionsgröße von 2, sodass ihre Ausgabe eine 0. Dimensionsgröße von 2 hat. Wenn der Batchverarbeitungsvorgang die endgültige Form von (3, 3, 4) aufteilt, wird (1, 3, 4) für Anfrage 1 und (2, 3, 4) für Anfrage 2 erzeugt.

Wenn diese Anforderung nicht erfüllt wird, führt das zu Fehlern wie: Batching input tensors supplied in a given op invocation must have equal 0th-dimension size.
Die Größe der 0. Dimension der Form jedes Ausgabetensors muss die Summe der 0. Dimensionsgröße des Eingabetensors sein (zuzüglich des Paddings, das durch den Batchverarbeitungsvorgang eingeführt wurde, um die nächstgrößere allowed_batch_size zu erreichen). Dadurch kann der Batchverarbeitungsvorgang die Ausgabetensoren entlang ihrer 0. Dimension basierend auf der 0. Dimension der Eingabetensoren aufteilen.

In der obigen Schritt-für-Schritt-Anleitung haben die Eingabetensoren die 0. Dimension aus Anfrage 1 und 2 aus Anfrage 2. Daher muss jeder Ausgabetensor eine 0. Dimension von 3 haben, da 1+2=3 ist. Der Ausgabetensor (3, 3, 4) erfüllt diese Anforderung. Wenn 3 keine gültige Batchgröße gewesen wäre, 4 aber keine gültige Batchgröße, hätte der Batchverarbeitungsvorgang die 0. Dimension der Eingaben von 3 auf 4 auffüllen mussten. In diesem Fall müsste jeder Ausgabetensor eine 0. Dimensionsgröße von 4 haben.

Wenn diese Anforderung nicht erfüllt wird, führt dies zu einem Fehler wie: Batched output tensor's 0th dimension does not equal the sum of the 0th dimension sizes of the input tensors.

Fehler bei Formanforderung beheben

Wenn eine Funktion in Batches zusammengefasst wird, muss die Form der „input_signature“ von @tf.function in der 0. Dimension (auch als Batchdimension bezeichnet) None haben. Wenn eine Signatur im Batch verarbeitet wird, achten Sie darauf, dass alle Eingaben in der 0. Dimension -1 enthalten.

Der BatchFunction-Vorgang unterstützt SparseTensors nicht als Ein- oder Ausgaben. Intern wird jeder dünnbesetzte Tensor als drei separate Tensoren dargestellt, die unterschiedliche Größen der 0. Dimension haben können.