Mantieni tutto organizzato con le raccolte Salva e classifica i contenuti in base alle tue preferenze.

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

  1. Crea un file denominato Dockerfile nella directory attuale e incolla il seguente testo

    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. Crea bucket Cloud Storage

    gsutil mb -c standard -l europe-west4 gs://your-bucket-name

  3. 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

  4. Copia il Dockerfile nella tua VM TPU

    gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:

  5. SSH nella VM TPU

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

  6. Crea l'immagine Docker

    sudo docker build -t your-image-name .

  7. Avvia il container Docker

    sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash

  8. 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

  9. 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

  1. Crea un file denominato Dockerfile nella directory attuale e incolla il seguente testo

    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. 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

  3. Copia il Dockerfile nella tua VM TPU

    gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:

  4. SSH nella VM TPU

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

  5. Crea l'immagine Docker

    sudo docker build -t your-image-name .

  6. Avvia un container Docker

    sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash

  7. 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.

  1. Digita exit per uscire dal container Docker
  2. Digita exit per uscire dalla VM TPU
  3. 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

  1. 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

  2. SSH nella VM TPU

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

  3. 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

  4. Installa libtpu

    pip install torch_xla[tpuvm]
export XRT_TPU_CONFIG="localservice;0;localhost:51011"

  5. Clona il repository PyTorch XLA

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

  6. 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.

  1. Digita exit per uscire dal container Docker
  2. Digita exit per uscire dalla VM TPU
  3. 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

  1. 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

  2. SSH nella VM TPU

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

  3. Avvia daemon Docker in VM TPU

    sudo systemctl start docker

  4. Avvia container Docker

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

  5. Installa JAX

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

  6. Installa FLAX

    pip install --upgrade clu
git clone https://github.com/google/flax.git
pip install --user -e flax

  7. 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.

  1. Digita exit per uscire dal container Docker
  2. Digita exit per uscire dalla VM TPU

  3. 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.

  1. 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 utente
    • your-project-name con il nome del tuo progetto Google Cloud
    • your-repo-name con il nome del tuo repository Artifact Registry
    • your-image-name con un nome per l'immagine Docker
    • your-tag con un tag per l'immagine docker
    • your-tpu-name con un nome per la TPU

  2. Crea un file denominato Dockerfile nella directory attuale e incolla il seguente testo

    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. Crea l'immagine Docker

    docker build -t $IMAGE .

  4. 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

  5. Esegui il push dell'immagine Docker ad Artifact Registry

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

  6. Crea una VM TPU

    gcloud compute tpus tpu-vm create $TPU_NAME \
--zone $ZONE \
--accelerator-type $ACCEL_TYPE \
--version tpu-vm-base

  7. 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"

  8. 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"

  9. 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