Esegui le applicazioni Cloud TPU in un container Docker
I container Docker semplificano la configurazione delle applicazioni combinando il codice e tutte le dipendenze necessarie in un unico pacchetto distribuibile. Puoi eseguire i container Docker all'interno delle VM TPU per semplificare la configurazione e la condivisione delle tue applicazioni Cloud TPU. Questo documento descrive come configurare un container Docker per ogni framework ML supportato da Cloud TPU.
Addestra un modello TensorFlow in un container Docker
Dispositivo TPU
Crea un file denominato
Dockerfile
nella directory attuale e incolla il seguente testoFROM 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
Crea bucket Cloud Storage
gsutil mb -c standard -l europe-west4 gs://your-bucket-name
Crea una VM TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-tf-2.11.0
Copia il Dockerfile nella tua VM TPU
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
SSH nella VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Crea l'immagine Docker
sudo docker build -t your-image-name .
Avvia il container Docker
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
Imposta le variabili di ambiente
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
Addestra ResNet
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"
pod di TPU
Crea un file denominato
Dockerfile
nella directory attuale e incolla il seguente testoFROM 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
Crea una VM TPU
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
Copia il Dockerfile nella tua VM TPU
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
SSH nella VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Crea l'immagine Docker
sudo docker build -t your-image-name .
Avvia un container Docker
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
Addestra ResNet
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"
Al termine dello script di addestramento, assicurati di eseguire la pulizia delle risorse.
- Digita
exit
per uscire dal container Docker - Digita
exit
per uscire dalla VM TPU - Elimina la VM TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Addestra un modello PyTorch in un container Docker
Crea VM di Cloud TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-pt-1.13
SSH nella VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Avvia un container nella VM TPU
sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm bash
Installa libtpu
pip install torch_xla[tpuvm] export XRT_TPU_CONFIG="localservice;0;localhost:51011"
Clona il repository PyTorch XLA
git clone --recursive https://github.com/pytorch/xla.git
Addestra ImageNet
python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
Al termine dello script di addestramento, assicurati di eseguire la pulizia delle risorse.
- Digita
exit
per uscire dal container Docker - Digita
exit
per uscire dalla VM TPU - Elimina la VM TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Addestra un modello JAX in un container Docker
Dispositivo TPU
Crea la VM TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-base
SSH nella VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Avvia daemon Docker in VM TPU
sudo systemctl start docker
Avvia container Docker
sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
Installa JAX
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Installa FLAX
pip install --upgrade clu git clone https://github.com/google/flax.git pip install --user -e flax
Esegui lo script di addestramento FLAX MNIST
cd flax/examples/mnist python3 main.py --workdir=/tmp/mnist \ --config=configs/default.py \ --config.learning_rate=0.05 \ --config.num_epochs=5
Al termine dello script di addestramento, assicurati di eseguire la pulizia delle risorse.
- Digita
exit
per uscire dal container Docker - Digita
exit
per uscire dalla VM TPU Elimina la VM TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
pod di TPU
Quando esegui il codice JAX su un pod TPU, devi eseguire contemporaneamente il codice JAX su tutti i worker di TPU. Un modo per farlo è utilizzare il comando gcloud compute tpus tpu-vm ssh
con i flag --worker=all
e --command
. La seguente procedura mostra come creare un'immagine Docker per semplificare la configurazione di ogni worker TPU.
Da un prompt del terminale, imposta le seguenti variabili di ambiente.
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
Sostituisci:
your-user-id
con il tuo ID utenteyour-project-name
con il nome del tuo progetto Google Cloudyour-repo-name
con il nome del tuo repository Artifact Registryyour-image-name
con un nome per l'immagine Dockeryour-tag
con un tag per l'immagine dockeryour-tpu-name
con un nome per la TPU
Crea un file denominato
Dockerfile
nella directory attuale e incolla il seguente testoFROM 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
Crea l'immagine Docker
docker build -t $IMAGE .
Aggiungi un tag alla tua immagine Docker prima di eseguirne il push su Artifact Registry. Per ulteriori informazioni sull'utilizzo di Artifact Registry, consulta Utilizzo delle immagini container.
docker tag $IMAGE_NAME europe-west-docker.pkg.dev/$PROJECT_NAME/$REPO_NAME/$IMAGE_NAME:$TAG
Esegui il push dell'immagine Docker ad Artifact Registry
docker push europe-west4-docker.pkg.dev/ml-writers/$REPO_NAME/$IMAGE_NAME:$TAG
Crea una VM TPU
gcloud compute tpus tpu-vm create $TPU_NAME \ --zone $ZONE \ --accelerator-type $ACCEL_TYPE \ --version tpu-vm-base
Esegui il pull dell'immagine Docker da Artifact Registry su tutti i worker TPU.
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"
Eseguire il container su tutti i worker di TPU.
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"
Esegui lo script di addestramento su tutti i worker di TPU:
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"
Al termine dello script di addestramento, assicurati di eseguire la pulizia delle risorse. Elimina la VM TPU con il seguente comando:
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Passaggi successivi
- Tutorial su Cloud TPU
- Guida dell'utente di Cloud TPU
- Architettura di sistema di Cloud TPU
- Eseguire il codice TensorFlow sulle sezioni di pod TPU
- Eseguire il codice JAX sulle sezioni dei pod TPU