Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.

Executar aplicativos do Cloud TPU em um contêiner do Docker

Os contêineres do Docker facilitam a configuração de aplicativos, combinando seu código e todas as dependências necessárias em um pacote distribuível. É possível executar contêineres do Docker em VMs de TPU para simplificar a configuração e o compartilhamento de aplicativos do Cloud TPU. Neste documento, descrevemos como configurar um contêiner do Docker para cada framework de ML com suporte do Cloud TPU.

Treinar um modelo do TensorFlow em um contêiner do Docker

Dispositivo TPU

  1. Crie um arquivo chamado Dockerfile no diretório atual e cole o texto a seguir

    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. Criação do bucket do Cloud Storage

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

  3. Criar uma VM de 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. Copie o Dockerfile para sua VM da TPU

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

  5. Conectar-se por SSH à VM de TPU

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

  6. Criar a imagem do Docker

    sudo docker build -t your-image-name .

  7. Iniciar o contêiner do Docker

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

  8. Defina as variáveis de 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. Treinar 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 de TPU

  1. Crie um arquivo chamado Dockerfile no diretório atual e cole o texto a seguir

    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. Criar uma VM de 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. Copie o Dockerfile para sua VM da TPU

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

  4. Conectar-se por SSH à VM de TPU

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

  5. Criar a imagem do Docker

    sudo docker build -t your-image-name .

  6. Iniciar um contêiner do Docker

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

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

Após a conclusão do script de treinamento, limpe os recursos.

  1. Digite exit para sair do contêiner do Docker
  2. Digite exit para sair da VM da TPU
  3. Excluir a VM da TPU 
      $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Treinar um modelo PyTorch em um contêiner do Docker

  1. Criar VM do 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 na VM da TPU

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

  3. Iniciar um contêiner na VM da TPU

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

  4. Instalar o libtpu

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

  5. Clonar o repositório do PyTorch XLA

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

  6. Treinar o ImageNet

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

Após a conclusão do script de treinamento, limpe os recursos.

  1. Digite exit para sair do contêiner do Docker
  2. Digite exit para sair da VM da TPU
  3. Excluir a VM da TPU 
    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Treinar um modelo JAX em um contêiner do Docker

Dispositivo TPU

  1. Criar a VM de TPU

    gcloud compute tpus tpu-vm create your-tpu-name \
--zone=europe-west4-a \
--accelerator-type=v2-8 \
--version=tpu-vm-base

  2. SSH na VM da TPU

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

  3. Iniciar o daemon do Docker na VM da TPU

    sudo systemctl start docker

  4. Iniciar contêiner do Docker

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

  5. Instalar o JAX

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

  6. Instalar o MXX

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

  7. Execute o script de treinamento MNIST do FLAX

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

Após a conclusão do script de treinamento, limpe os recursos.

  1. Digite exit para sair do contêiner do Docker
  2. Digite exit para sair da VM da TPU

  3. Excluir a VM da TPU

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

Pod de TPU

Ao executar o código JAX em um pod de TPU, é preciso executar esse código em todos os workers da TPU ao mesmo tempo. Uma maneira de fazer isso é usar o comando gcloud compute tpus tpu-vm ssh com as sinalizações --worker=all e --command. O procedimento a seguir mostra como criar uma imagem do Docker para facilitar a configuração de cada worker de TPU.

  1. Em um prompt do terminal, defina as variáveis de ambiente a seguir. 

    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

    Substitua:

    • your-user-id pelo seu ID de usuário
    • your-project-name pelo nome do projeto do Google Cloud
    • your-repo-name pelo nome do seu repositório do Artifact Registry
    • your-image-name por um nome para sua imagem do Docker
    • your-tag por uma tag para a imagem do Docker
    • your-tpu-name por um nome para sua TPU

  2. Crie um arquivo chamado Dockerfile no diretório atual e cole o texto a seguir

    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. Criar a imagem do Docker

    docker build -t $IMAGE .

  4. Adicione uma tag à sua imagem do Docker antes de enviá-la ao Artifact Registry. Para saber mais sobre como trabalhar com o Artifact Registry, consulte Trabalhar com imagens de contêiner.

    docker tag $IMAGE_NAME europe-west-docker.pkg.dev/$PROJECT_NAME/$REPO_NAME/$IMAGE_NAME:$TAG

  5. Enviar a imagem do Docker para o Artifact Registry

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

  6. Criar uma VM de TPU

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

  7. Extraia a imagem do Docker do Artifact Registry em todos os workers da 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. Execute o contêiner em todos os workers da 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. Execute o script de treinamento em todos os workers da 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"

Após a conclusão do script de treinamento, limpe os recursos. Exclua a VM da TPU usando o seguinte comando:

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

A seguir