Auf TPU-Pods trainieren

Überblick

Einer der Vorteile von Cloud TPUs ist, dass Sie das Training zwischen TPU-Typen verschiedener Größen skalieren können. Alle TPU-Typen verwenden die gleiche datenparallele Architektur. Der einzige Unterschied besteht darin, dass sich die Parallelität von 8 Kernen auf 2.048 Kerne erhöht.

Um eine größere Anzahl von TPUs optimal nutzen zu können, müssen Sie mehrere Trainingsaufgabenparameter optimieren. In diesem Dokument werden einige häufig auftretende Probleme und die Änderungen beschrieben, die Sie dann an Ihren Modellen vornehmen müssen.

Batchgröße und Trainingsschritte skalieren

Wenn Sie bei größeren TPU-Typen eine lineare Skalierung erzielen möchten, sollten Sie die Batchgröße pro Kern beibehalten.

Wenn Sie beispielsweise für einen v2-8-TPU-Typ eine Batchgröße von 1.024 verwenden, sollten Sie für einen v2-32-TPU-Typ eine Batchgröße von 4.096 (4 * 1.024) verwenden. Dadurch wird die TPU-Hardware bestmöglich genutzt. Sie können kleinere Batchgrößen verwenden, aber das Training wird dann nicht linear skaliert.

Viele Modelle enthalten ein Flag train_steps, wobei ein Schritt einem einzelnen Datenbatch entspricht. Wenn Sie die Batchgröße erhöhen, verkleinern Sie die Anzahl der Trainingsschritte, sodass die Gesamtzahl der Trainingsbeispiele gleich bleibt.

Wenn die ursprüngliche Konfiguration beispielsweise eine Batchgröße von 1024 für 1.000 Schritte hatte, können Sie die Batchgröße auf 4096 erhöhen und die Schritte auf 250 (1000 / 4) reduzieren. Wenn Ihr Modell das Flag epochs verwendet, müssen Sie die Anzahl der Schritte nicht skalieren.

Größere Batches können zu einem geänderten Konvergenzverhalten des Modells führen. Es empfiehlt sich daher ggf., auch gewisse Hyperparameter abzustimmen.

Optimierungstools für große Batches verwenden

Wenn Sie auf größere Batches hochskalieren, müssen Sie möglicherweise auch ein Optimierungstool nutzen, dass große Batches unterstützt. Beispielsweise wird beim Referenzmodell Resnet für größere TPU-Typen der LARSOptimizer verwendet.

Regionale Google Cloud Storage-Buckets in der Region des TPU-Pods verwenden

Die Best Practice für das TPU-Training besteht im Allgemeinen darin, immer nur Ressourcen aus einer Region zu verwenden. Die Region der Ressourcen ist dann besonders wichtig, wenn Sie TPU-Pods verwenden, weil die Rate der Datenübertragung von Google Cloud Storage dann meist höher ist. Verwenden Sie für Training-Datasets und Prüfpunkte unbedingt einen regionalen Google Cloud Storage-Bucket in derselben Region wie die TPU.

TPU-Name und -Region für TPUClusterResolver bereitstellen

Wenn Sie für das Training einen TPU-Typ mit einem einzelnen Gerät verwenden, können Sie den TPU-Namen oder eine IP-Adresse angeben, zum Beispiel grpc://1.2.3.4:8470.

Bei TPU-Pods müssen Sie den TPU-Namen verwenden, damit TensorFlow die IP-Adressen aller für die Trainingsdistribution verfügbaren Hosts ermitteln kann.

Wenn Sie für einen TPU-Pod die IP-Adresse statt dem Namen angeben, wird der folgende Fehler ausgegeben:

ValueError: TPUConfig.num_shards is not set correctly ....

Wenn Sie Ihr Modell und Ihre TPU in verschiedenen Zonen oder Regionen ausführen müssen, müssen Sie die TPU-Zone für TPUClusterResolver angeben.

Sie können die Flags tpu und tpu_zone an Referenzmodelle übergeben oder die Werte direkt für TPUClusterResolver festlegen. Beispiel:

cluster = tf.contrib.cluster_resolver.TPUClusterResolver(
  tpu='tpuname',
  zone='europe-west4-a',
  project='myproject')

Wenn Sie TPUEstimator verwenden, kann es auch vorkommen, dass der Cluster-Resolver nicht an tf.contrib.tpu.RunConfig übergeben wird. Das kann passieren, wenn RunConfig mit dem Parameter master initialisiert wurde. Initialisieren Sie es stattdessen mit dem Parameter cluster.

config = tf.contrib.tpu.RunConfig(cluster=cluster)

Wenn Sie einen TPU-Pod und tf.Session ohne TPUEstimator verwenden, müssen Sie die Sitzungskonfiguration so konfigurieren:

with tf.Session(cluster.get_master(), config=config.session_config) as sess:

Evaluierung auf TPU-Pods

Das Evaluieren wird auf TPU-Pods nicht unterstützt und ist dort auch nicht kosteneffizient. Google empfiehlt, Evaluierungen auf einem TPU-Typ mit einem einzelnen Gerät auszuführen und dafür das Flag mode zu verwenden, das von den meisten Modellen unterstützt wird.

Für das Training auf einem TPU-Pod oder einem TPU-Typ mit einem einzelnen Gerät legen Sie das Flag auf train fest. Für das Evaluieren legen Sie eval ausschließlich für TPU-Typen mit einem einzelnen Gerät fest. Auf TPU-Typen mit einem einzelnen Gerät können Sie auch den Modus train_and_eval verwenden, jedoch nicht auf TPU-Pods.