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
Erstellen Sie in Ihrem aktuellen Verzeichnis eine Datei mit dem Namen
Dockerfile
und fügen Sie den folgenden Text einFROM 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
Cloud Storage-Bucket erstellen
gsutil mb -c standard -l europe-west4 gs://your-bucket-name
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
Dockerfile auf die TPU-VM kopieren
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
Stellen Sie eine SSH-Verbindung zur TPU-VM her:
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Docker-Image erstellen
sudo docker build -t your-image-name .
Docker-Container starten
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
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
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
Erstellen Sie in Ihrem aktuellen Verzeichnis eine Datei mit dem Namen
Dockerfile
und fügen Sie den folgenden Text einFROM 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
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
Dockerfile auf die TPU-VM kopieren
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
Stellen Sie eine SSH-Verbindung zur TPU-VM her:
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Docker-Image erstellen
sudo docker build -t your-image-name .
Docker-Container starten
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
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.
- Geben Sie
exit
ein, um den Docker-Container zu beenden - Geben Sie
exit
ein, um die TPU-VM zu beenden - TPU-VM löschen
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
PyTorch-Modell in einem Docker-Container trainieren
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
SSH-Verbindung zur TPU-VM herstellen
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
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
libtpu installieren
pip install torch_xla[tpuvm] export XRT_TPU_CONFIG="localservice;0;localhost:51011"
PyTorch XLA-Repository klonen
git clone --recursive https://github.com/pytorch/xla.git
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.
- Geben Sie
exit
ein, um den Docker-Container zu beenden - Geben Sie
exit
ein, um die TPU-VM zu beenden - 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
TPU-VM erstellen
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-base
SSH-Verbindung zur TPU-VM herstellen
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Docker-Daemon in TPU-VM starten
sudo systemctl start docker
Docker-Container starten
sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
JAX installieren
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
FLAX installieren
pip install --upgrade clu git clone https://github.com/google/flax.git pip install --user -e flax
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.
- Geben Sie
exit
ein, um den Docker-Container zu beenden - Geben Sie
exit
ein, um die TPU-VM zu beenden 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.
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-IDyour-project-name
durch den Namen Ihres Google Cloud-Projektsyour-repo-name
durch den Namen Ihres Artifact Registry-Repositorysyour-image-name
durch einen Namen für Ihr Docker-Imageyour-tag
mit einem Tag für das Docker-Imageyour-tpu-name
durch einen Namen für Ihre TPU
Erstellen Sie in Ihrem aktuellen Verzeichnis eine Datei mit dem Namen
Dockerfile
und fügen Sie den folgenden Text einFROM 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
Docker-Image erstellen
docker build -t $IMAGE .
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
Docker-Image per Push an die Artifact Registry übertragen
docker push europe-west4-docker.pkg.dev/ml-writers/$REPO_NAME/$IMAGE_NAME:$TAG
TPU-VM erstellen
gcloud compute tpus tpu-vm create $TPU_NAME \ --zone $ZONE \ --accelerator-type $ACCEL_TYPE \ --version tpu-vm-base
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"
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"
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
- Cloud TPU-Anleitungen
- Cloud TPU-Nutzerhandbuch
- Cloud TPU-Systemarchitektur
- TensorFlow-Code auf TPU Pod-Slices ausführen
- JAX-Code auf TPU Pod-Slices ausführen