TPUs zum Trainieren eines Modells verwenden

Tensor Processing Units (TPUs) sind von Google speziell entwickelte ASICs, mit denen beim maschinellen Lernen entstehende Arbeitslasten beschleunigt werden. Mit Cloud TPU können Sie Ihre Trainingsjobs in AI Platform Training ausführen. AI Platform Training bietet eine Benutzeroberfläche für die Auftragsverwaltung, sodass Sie die TPU nicht selbst verwalten müssen. Stattdessen können Sie die AI Platform Training jobs API genauso verwenden, wie Sie sie für das Training auf einer CPU oder einer GPU verwenden.

Mit den übergeordneten TensorFlow APIs können Sie Ihre Modelle auf der Cloud-TPU-Hardware ausführen.

Google Cloud-Umgebung einrichten und testen

Konfigurieren Sie Ihre Google Cloud-Umgebung. Lesen Sie dafür im Startleitfaden den Abschnitt zur Einrichtung.

Cloud TPU Zugriff auf das Projekt gewähren

Führen Sie die folgenden Schritte aus, um den Namen des Cloud TPU-Dienstkontos zu autorisieren, der Ihrem Google Cloud-Projekt zugeordnet ist:

  1. Rufen Sie den Namen Ihres Cloud TPU-Dienstkontos durch einen Aufruf von projects.getConfig ab. Beispiel:

    curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
            https://ml.googleapis.com/v1/projects/<your-project-id>:getConfig
        
  2. Speichern Sie den Wert des Felds tpuServiceAccount, der von der API zurückgegeben wird.

Fügen Sie dem Projekt das Cloud TPU-Dienstkonto als Mitglied mit der Rolle Cloud ML-Dienst-Agent hinzu. Führen Sie die folgenden Schritte in der Google Cloud Console oder mit dem Befehl gcloud aus:

Console

  1. Melden Sie sich in der Google Cloud Console an und wählen Sie das Projekt aus, in dem Sie die TPU verwenden.
  2. Wählen Sie IAM und Verwaltung > IAM aus.
  3. Klicken Sie auf die Schaltfläche Hinzufügen, um ein Mitglied zum Projekt hinzuzufügen.
  4. Geben Sie das TPU-Dienstkonto in das Textfeld Mitglieder ein.
  5. Klicken Sie auf die Drop-down-Liste Rollen.
  6. Aktivieren Sie die Rolle Cloud ML-Dienst-Agent (Dienstverwaltung > Cloud ML-Dienst-Agent).

gcloud

  1. Legen Sie Umgebungsvariablen fest, die Ihre Projekt-ID und das Cloud TPU-Dienstkonto enthalten:

    PROJECT_ID=your-project-id
        SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
        
  2. Weisen Sie dem Cloud TPU-Dienstkonto die Rolle ml.serviceAgent zu:

    gcloud projects add-iam-policy-binding $PROJECT_ID \
            --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
        

Weitere Informationen über das Zuweisen von Rollen zu Dienstkonten finden Sie in der Cloud IAM-Dokumentation.

Beispielmodell ResNet-50 ausführen

In diesem Abschnitt wird gezeigt, wie Sie das TensorFlow-Referenzmodell ResNet-50 mit einem fiktiven Dataset trainieren. Dieses steht unter gs://cloud-tpu-test-datasets/fake_imagenet zur Verfügung. In diesem Beispieljob wird für die Maschinenkonfiguration die vordefinierte Skalierungsstufe BASIC_TPU verwendet. Weiter unten in dieser Anleitung erfahren Sie, wie Sie eine benutzerdefinierte Konfiguration einrichten.

Führen Sie die folgenden Befehle aus, um den Code abzurufen und Ihren Trainingsjob an AI Platform Training zu senden:

  1. Laden Sie den Code für das Referenzmodell herunter:

    mkdir tpu-demos && cd tpu-demos
        wget https://github.com/tensorflow/tpu/archive/r1.15.tar.gz
        tar -xzvf r1.15.tar.gz && rm r1.15.tar.gz
        
  2. Öffnen Sie in der entpackten Verzeichnisstruktur das Verzeichnis official:

    cd tpu-r1.15/models/official/
        
  3. Bearbeiten Sie ./resnet/resnet_main.py und ändern Sie den Code so, dass beim Importieren von untergeordneten Modulen explizite relative Importe verwendet werden. Ändern Sie zum Beispiel den Eintrag

    import resnet_model
        

    in

    from . import resnet_model
        

    Ändern Sie nach dem obigen Muster alle anderen Importe in der Datei und speichern Sie die Datei anschließend.

  4. Suchen Sie in anderen Dateien des Beispiels nach allen Importen untergeordneter Module und aktualisieren Sie sie, um auch dort explizite relative Importe zu verwenden.

  5. Richten Sie einige Umgebungsvariablen ein:

    JOB_NAME=tpu_1
        STAGING_BUCKET=gs://my_bucket_for_staging
        REGION=us-central1
        DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
        OUTPUT_PATH=gs://my_bucket_for_model_output
        

    Der Zugriff auf TPUs ist derzeit in folgenden Regionen möglich:

    • us-central1

  6. Senden Sie Ihren Trainingsjob mit dem Befehl gcloud ai-platform jobs submit training:

    gcloud ai-platform jobs submit training $JOB_NAME \
                --staging-bucket $STAGING_BUCKET \
                --runtime-version 1.15 \
                --scale-tier BASIC_TPU \
                --module-name resnet.resnet_main \
                --package-path resnet/ \
                --region $REGION \
                -- \
                --data_dir=$DATA_DIR \
                --model_dir=$OUTPUT_PATH
        

Weitere Möglichkeiten zum Trainieren von Modellen mit Cloud TPU

Im ersten Teil dieser Anleitung wurde beschrieben, wie Sie den ResNet-50-Beispielcode verwenden. In diesem Abschnitt erfahren Sie mehr über das Konfigurieren eines Jobs und das Trainieren eines Modells in AI Platform Training mit Cloud TPU.

Region angeben, die TPUs anbietet

Sie müssen Ihren Job in einer Region ausführen, in der TPUs verfügbar sind. Der Zugriff auf TPUs ist derzeit in folgenden Regionen möglich:

  • us-central1

Ausführliche Informationen zu den verfügbaren Regionen für AI Platform Training-Dienste, einschließlich Modelltraining und Online-/Batchvorhersagen, finden Sie im Leitfaden für Regionen.

Versionsverwaltung für TensorFlow und AI Platform Training

Die AI Platform Training-Laufzeitversionen 1.13, 1.14 und 1.15 sind zum Trainieren Ihrer Modelle in Cloud TPU verfügbar. Weitere Informationen zu AI Platform Training-Laufzeitversionen und den entsprechenden TensorFlow-Versionen.

Die Versionierungsrichtlinie entspricht der Richtlinie für Cloud TPU. Geben Sie in der Trainingsjobanfrage unbedingt eine Laufzeitversion an, die für TPUs verfügbar ist und der Version von TensorFlow im Trainingscode entspricht.

Mit dem TPU-gRPC-Server verbinden

Innerhalb Ihres TensorFlow-Programms sollten Sie mit TPUClusterResolver eine Verbindung zum TPU-gRPC-Server herstellen, der auf der TPU-VM ausgeführt wird. TPUClusterResolver liefert die IP-Adresse und den Port von Cloud TPU.

Im folgenden Beispiel ist zu sehen, wie im ResNet-50-Beispielcode TPUClusterResolver verwendet wird:

tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
        FLAGS.tpu,
        zone=FLAGS.tpu_zone,
        project=FLAGS.gcp_project)

    config = tpu_config.RunConfig(
        cluster=tpu_cluster_resolver,
        model_dir=FLAGS.model_dir,
        save_checkpoints_steps=max(600, FLAGS.iterations_per_loop),
        tpu_config=tpu_config.TPUConfig(
            iterations_per_loop=FLAGS.iterations_per_loop,
            num_shards=FLAGS.num_cores,
            per_host_input_for_training=tpu_config.InputPipelineConfig.PER_HOST_V2))  # pylint: disable=line-too-long

TPUs Vorgänge zuweisen

Verwenden Sie zur Nutzung der TPUs auf einer Maschine die TensorFlow TPUEstimator API, die von der übergeordneten TensorFlow Estimator API abgeleitet wird.

  • TPUEstimator verarbeitet viele Details der Ausführung auf TPU-Geräten, z. B. das Replizieren von Eingaben und Modellen für jeden Kern und die regelmäßige Rückgabe an den Host, um Hooks auszuführen.
  • Die übergeordnete TensorFlow API bietet auch viele weitere praktische Funktionen. Insbesondere übernimmt die API die Speicherung und Wiederherstellung der Modellkontrollpunkte, damit Sie einen unterbrochenen Trainingsjob an dem Punkt fortsetzen können, an dem er angehalten wurde.

Hier finden Sie die Liste der TensorFlow-Vorgänge, die in Cloud TPU verfügbar sind.

Benutzerdefinierte TPU-Maschine konfigurieren

Ein TPU-Trainingsjob wird in einer Konfiguration mit zwei VMs ausgeführt. Eine VM (der Master) führt Ihren Python-Code aus. Der Master steuert den TensorFlow-Server, der auf einem TPU-Worker ausgeführt wird.

Wenn Sie eine TPU mit AI Platform Training verwenden möchten, haben Sie drei Möglichkeiten, den Trainingsjob für den Zugriff auf eine TPU-fähige Maschine zu konfigurieren:

  • Verwenden Sie die Skalierungsstufe BASIC_TPU. Sie können diese Methode für den Zugriff auf TPU v2-Beschleuniger verwenden.
  • Verwenden Sie für die Master-VM einen cloud_tpu-Worker und einen Legacy-Maschinentyp. Sie können diese Methode für den Zugriff auf TPU v2-Beschleuniger verwenden.
  • Verwenden Sie für die Master-VM einen cloud_tpu-Worker und einen Compute Engine-Maschinentyp. Sie können diese Methode für den Zugriff auf TPU v2- oder TPU v3-Beschleuniger verwenden. TPU v3-Beschleuniger sind in der Betaversion verfügbar.

Einfache TPU-fähige Maschine

Setzen Sie die Skalierungsstufe auf BASIC_TPU, um eine Master-VM und eine TPU-VM mit einer einzigen TPU und acht TPU v2-Kernen zu erhalten, wie Sie es beim Ausführen des vorherigen Beispiels getan haben.

TPU-Worker in einer Konfiguration mit Legacy-Maschinentyp

Wenn Sie auf der Master-VM mehr Rechenressourcen benötigen, können Sie alternativ eine benutzerdefinierte Maschinenkonfiguration einrichten:

  • Setzen Sie die Skalierungsstufe auf CUSTOM.
  • Konfigurieren Sie die Master-VM so, dass ein Legacy-Maschinentyp entsprechend den Jobanforderungen verwendet wird.
  • Setzen Sie workerType auf cloud_tpu, damit Sie eine TPU-VM mit einer Cloud TPU und acht TPU v2-Kernen erhalten.
  • Setzen Sie workerCount auf 1.
  • Geben Sie keinen Parameterserver an, wenn Sie eine Cloud TPU verwenden. Der Dienst weist die Jobanfrage zurück, wenn parameterServerCount größer Null ist.

Das folgende Beispiel zeigt eine config.yaml-Datei, die diesen Konfigurationstyp verwendet:

trainingInput:
      scaleTier: CUSTOM
      masterType: complex_model_m
      workerType: cloud_tpu
      workerCount: 1
    

TPU-Worker in einer Konfiguration mit Compute Engine-Maschinentyp

Sie können auch eine benutzerdefinierte Maschinenkonfiguration mit einem Compute Engine-Maschinentyp für die Master-VM einrichten und an die TPU-VM eine acceleratorConfig anhängen.

Mit dieser Art der Konfiguration können Sie einen TPU-Worker mit acht TPU v2-Kernen (ähnlich wie bei einer Konfiguration ohne acceleratorConfig) oder einen TPU-Worker mit acht TPU v3-Kernen (Beta) einrichten. Hier erfahren Sie mehr über den Unterschied zwischen TPU v2- und TPU v3-Beschleunigern.

Die Verwendung eines Compute Engine-Maschinentyps bietet auch mehr Flexibilität bei der Konfiguration der Master-VM:

  • Setzen Sie die Skalierungsstufe auf CUSTOM.
  • Konfigurieren Sie die Master-VM für die Verwendung eines Compute Engine-Maschinentyps, der Ihren Jobanforderungen entspricht.
  • Setzen Sie workerType auf cloud_tpu.
  • Fügen Sie eine workerConfig mit dem Feld acceleratorConfig hinzu. In der acceleratorConfig setzen Sie type auf TPU_V2 oder TPU_V3 sowie count auf 8. Sie dürfen keine andere Anzahl von TPU-Kernen anhängen.
  • Setzen Sie workerCount auf 1.
  • Geben Sie keinen Parameterserver an, wenn Sie eine Cloud TPU verwenden. Der Dienst weist die Jobanfrage zurück, wenn parameterServerCount größer Null ist.

Das folgende Beispiel zeigt eine config.yaml-Datei, die diesen Konfigurationstyp verwendet:

TPU v2

trainingInput:
      scaleTier: CUSTOM
      masterType: n1-highcpu-16
      workerType: cloud_tpu
      workerCount: 1
      workerConfig:
        acceleratorConfig:
          type: TPU_V2
          count: 8
    

TPU v3 (Beta)

trainingInput:
      scaleTier: CUSTOM
      masterType: n1-highcpu-16
      workerType: cloud_tpu
      workerCount: 1
      workerConfig:
        acceleratorConfig:
          type: TPU_V3
          count: 8
    

Benutzerdefinierte Container für einen TPU-Worker verwenden

Wenn Sie einen benutzerdefinierten Container auf Ihrem TPU-Worker ausführen möchten, anstatt eine der AI Platform Training-Laufzeitversionen zu verwenden, die TPUs unterstützen, müssen Sie beim Senden Ihres Trainingsjobs ein zusätzliches Konfigurationsfeld angeben. Legen Sie die tpuTfVersion auf eine Laufzeitversion fest, die die von Ihrem Container verwendete TensorFlow-Version enthält. Sie müssen eine Laufzeitversion angeben, die derzeit für das Training mit TPUs unterstützt wird.

Da Sie Ihren Job für die Verwendung eines benutzerdefinierten Containers konfigurieren, verwendet AI Platform Training bei der Ausführung Ihres Trainingsjobs nicht die Umgebung dieser Laufzeitversion. Für AI Platform Training ist dieses Feld jedoch erforderlich, damit der TPU-Worker ordnungsgemäß auf die Version von TensorFlow vorbereitet werden kann, die in Ihrem benutzerdefinierten Container verwendet wird.

Das folgende Beispiel zeigt eine config.yaml-Datei mit einer ähnlichen TPU-Konfiguration wie im vorherigen Abschnitt. In diesem Fall führen die Master-VM und der TPU-Worker jeweils unterschiedliche benutzerdefinierte Container aus:

TPU v2

trainingInput:
      scaleTier: CUSTOM
      masterType: n1-highcpu-16
      masterConfig:
        imageUri: gcr.io/your-project-id/your-master-image-name:your-master-tag-name
      workerType: cloud_tpu
      workerCount: 1
      workerConfig:
        imageUri: gcr.io/your-project-id/your-worker-image-name:your-worker-tag-name
        tpuTfVersion: 1.15
        acceleratorConfig:
          type: TPU_V2
          count: 8
    

TPU v3 (Beta)

trainingInput:
      scaleTier: CUSTOM
      masterType: n1-highcpu-16
      masterConfig:
        imageUri: gcr.io/your-project-id/your-master-image-name:your-master-tag-name
      workerType: cloud_tpu
      workerCount: 1
      workerConfig:
        imageUri: gcr.io/your-project-id/your-worker-image-name:your-worker-tag-name
        tpuTfVersion: 1.15
        acceleratorConfig:
          type: TPU_V3
          count: 8
    

Wenn Sie den Trainingsjob mit dem Befehl gcloud beta ai-platform jobs submit training senden, können Sie das API-Feld tpuTfVersion mit dem Flag --tpu-tf-version statt in einer config.yaml-Datei angeben.

Weitere Informationen zum verteilten Training mit benutzerdefinierten Containern

Nächste Schritte