Mit Sammlungen den Überblick behalten Sie können Inhalte basierend auf Ihren Einstellungen speichern und kategorisieren.

Cloud TPU-Anwendungen in einem Docker-Container ausführen

Docker-Container erleichtern die Konfiguration von Anwendungen, indem sie Ihren Code und alle benötigten Abhängigkeiten in einem verteilbaren Paket kombinieren. Sie können Docker-Container in TPU-VMs ausführen, um die Konfiguration und Freigabe Ihrer Cloud TPU-Anwendungen zu vereinfachen. In diesem Dokument wird beschrieben, wie Sie einen Docker-Container für jedes von Cloud TPU unterstützte ML-Framework einrichten.

TensorFlow-Modell in einem Docker-Container trainieren

TPU-Gerät

  1. Erstellen Sie in Ihrem aktuellen Verzeichnis eine Datei mit dem Namen Dockerfile und fügen Sie den folgenden Text ein

    FROM python:3.8
    RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.10.0/tensorflow-2.10.0-cp38-cp38-linux_x86_64.whl
    RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.4.0/libtpu.so -o /lib/libtpu.so
    RUN git clone https://github.com/tensorflow/models.git
    WORKDIR ./models
    RUN pip install -r official/requirements.txt
    ENV PYTHONPATH=/models
    
  2. Cloud Storage-Bucket erstellen

    gsutil mb -c standard -l europe-west4 gs://your-bucket-name
    
  3. TPU-VM erstellen

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v2-8 \
    --version=tpu-vm-tf-2.11.0
    
  4. Dockerfile auf die TPU-VM kopieren

    gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
    
  5. Stellen Sie eine SSH-Verbindung zur TPU-VM her:

    gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
    
  6. Docker-Image erstellen

    sudo docker build -t your-image-name .
    
  7. Docker-Container starten

    sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
    
  8. Umgebungsvariablen festlegen

    export STORAGE_BUCKET=gs://your-bucket-name
    export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
    export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x
    export ZONE=europe-west4-a
    export TPU_NAME=local
    
  9. ResNet trainieren

    python3 official/vision/train.py \
    --tpu=local \
    --experiment=resnet_imagenet \
    --mode=train_and_eval \
    --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
    --model_dir=${MODEL_DIR} \
    --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,trainer.train_steps=100"
    

TPU-Pod

  1. Erstellen Sie in Ihrem aktuellen Verzeichnis eine Datei mit dem Namen Dockerfile und fügen Sie den folgenden Text ein

    FROM python:3.8
    RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.10.0/tensorflow-2.10.0-cp38-cp38-linux_x86_64.whl
    RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.4.0/libtpu.so -o /lib/libtpu.so
    RUN git clone https://github.com/tensorflow/models.git
    WORKDIR ./models
    RUN pip install -r official/requirements.txt
    ENV PYTHONPATH=/models
    
  2. TPU-VM erstellen

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v3-32 \
    --version=tpu-vm-tf-2.11.0-pod
    
  3. Dockerfile auf die TPU-VM kopieren

    gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
    
  4. Stellen Sie eine SSH-Verbindung zur TPU-VM her:

    gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
    
  5. Docker-Image erstellen

    sudo docker build -t your-image-name .
    
  6. Docker-Container starten

    sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
    
  7. ResNet trainieren

    python3 official/vision/train.py \
    --tpu=your-tpu-name \
    --experiment=resnet_imagenet \
    --mode=train_and_eval \
    --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
    --model_dir=${MODEL_DIR} \
    --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
    

Bereinigen Sie nach Abschluss des Trainingsskripts die Ressourcen.

  1. Geben Sie exit ein, um den Docker-Container zu beenden
  2. Geben Sie exit ein, um die TPU-VM zu beenden
  3. TPU-VM löschen
      $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
    

PyTorch-Modell in einem Docker-Container trainieren

  1. Cloud TPU-VM erstellen

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v2-8 \
    --version=tpu-vm-pt-1.13
    
  2. SSH-Verbindung zur TPU-VM herstellen

    gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
    
  3. Container in der TPU-VM starten

    sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm bash
    
  4. libtpu installieren

    pip install torch_xla[tpuvm]
    export XRT_TPU_CONFIG="localservice;0;localhost:51011"
    
  5. PyTorch XLA-Repository klonen

    git clone --recursive https://github.com/pytorch/xla.git
    
  6. ImageNet trainieren

    python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
    

Bereinigen Sie nach Abschluss des Trainingsskripts die Ressourcen.

  1. Geben Sie exit ein, um den Docker-Container zu beenden
  2. Geben Sie exit ein, um die TPU-VM zu beenden
  3. TPU-VM löschen
    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
    

JAX-Modell in einem Docker-Container trainieren

TPU-Gerät

  1. TPU-VM erstellen

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v2-8 \
    --version=tpu-vm-base
    
  2. SSH-Verbindung zur TPU-VM herstellen

    gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a
    
  3. Docker-Daemon in TPU-VM starten

    sudo systemctl start docker
    
  4. Docker-Container starten

    sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
    
  5. JAX installieren

    pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    
  6. FLAX installieren

    pip install --upgrade clu
    git clone https://github.com/google/flax.git
    pip install --user -e flax
    
  7. FLAX MNIST-Trainingsskript ausführen

    cd flax/examples/mnist
    python3 main.py --workdir=/tmp/mnist \
    --config=configs/default.py \
    --config.learning_rate=0.05 \
    --config.num_epochs=5
    

Bereinigen Sie nach Abschluss des Trainingsskripts die Ressourcen.

  1. Geben Sie exit ein, um den Docker-Container zu beenden
  2. Geben Sie exit ein, um die TPU-VM zu beenden
  3. TPU-VM löschen

    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
    

TPU-Pod

Wenn Sie JAX-Code auf einem TPU-Pod ausführen, müssen Sie ihn auf allen TPU-Workern gleichzeitig ausführen. Eine Möglichkeit dazu ist, den Befehl gcloud compute tpus tpu-vm ssh mit den Flags --worker=all und --command zu verwenden. Das folgende Verfahren zeigt, wie Sie ein Docker-Image erstellen, um die Einrichtung der einzelnen TPU-Worker zu vereinfachen.

  1. Legen Sie in einer Terminalaufforderung die folgenden Umgebungsvariablen fest.

    export USER=your-user-id
    export PROJECT=your-project-name
    export REPO_NAME=your-repo-name
    export IMAGE_NAME=your-image-name
    export TAG=your-tag
    export TPU_NAME=your-tpu-name
    export ZONE=europe-west4-a
    export ACCEL_TYPE=v2-8
    

    Ersetzen Sie:

    • your-user-id durch Ihre Nutzer-ID
    • your-project-name durch den Namen Ihres Google Cloud-Projekts
    • your-repo-name durch den Namen Ihres Artifact Registry-Repositorys
    • your-image-name durch einen Namen für Ihr Docker-Image
    • your-tag mit einem Tag für das Docker-Image
    • your-tpu-name durch einen Namen für Ihre TPU
  2. Erstellen Sie in Ihrem aktuellen Verzeichnis eine Datei mit dem Namen Dockerfile und fügen Sie den folgenden Text ein

    FROM google-cloud-cli:latest
    RUN pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    RUN pip install --upgrade clu
    RUN git clone https://github.com/google/flax.git
    RUN pip install --user -e flax
    WORKDIR ./flax/examples/mnist
    
  3. Docker-Image erstellen

    docker build -t $IMAGE .
    
  4. Fügen Sie Ihrem Docker-Image ein Tag hinzu, bevor Sie es in die Artifact Registry übertragen. Weitere Informationen zum Arbeiten mit Artifact Registry finden Sie unter Mit Container-Images arbeiten.

    docker tag $IMAGE_NAME europe-west-docker.pkg.dev/$PROJECT_NAME/$REPO_NAME/$IMAGE_NAME:$TAG
    
  5. Docker-Image per Push an die Artifact Registry übertragen

    docker push europe-west4-docker.pkg.dev/ml-writers/$REPO_NAME/$IMAGE_NAME:$TAG
    
  6. TPU-VM erstellen

    gcloud compute tpus tpu-vm create $TPU_NAME \
    --zone $ZONE \
    --accelerator-type $ACCEL_TYPE \
    --version tpu-vm-base
    
  7. Rufen Sie das Docker-Image aus der Artifact Registry auf allen TPU-Workern ab.

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
    --command="sudo usermod -a -G docker $USER"
    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
    --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
    --command="docker pull europe-west4-docker.pkg.dev/your-project-name/your-repo-name/your-image-name:your-tag"
    
  8. Container auf allen TPU-Workern ausführen

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
    --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project-name/your-repo-name/your-image-name:your-tag bash"
    
  9. Führen Sie das Trainingsskript auf allen TPU-Workern aus:

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
    --zone $ZONE \
    --command="docker exec --privileged $CONTAINER_NAME python3 main.py --workdir=/tmp/mnist \
    --config=configs/default.py \
    --config.learning_rate=0.05 \
    --config.num_epochs=5"
    

Bereinigen Sie nach Abschluss des Trainingsskripts die Ressourcen. Löschen Sie die TPU-VM mit dem folgenden Befehl:

  $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
  

Nächste Schritte