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
Crie um arquivo chamado
Dockerfile
no diretório atual e cole o texto a seguirFROM 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
Criação do bucket do Cloud Storage
gsutil mb -c standard -l europe-west4 gs://your-bucket-name
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
Copie o Dockerfile para sua VM da TPU
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
Conectar-se por SSH à VM de TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Criar a imagem do Docker
sudo docker build -t your-image-name .
Iniciar o contêiner do Docker
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
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
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
Crie um arquivo chamado
Dockerfile
no diretório atual e cole o texto a seguirFROM 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
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
Copie o Dockerfile para sua VM da TPU
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
Conectar-se por SSH à VM de TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Criar a imagem do Docker
sudo docker build -t your-image-name .
Iniciar um contêiner do Docker
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
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.
- Digite
exit
para sair do contêiner do Docker - Digite
exit
para sair da VM da TPU - 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
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
SSH na VM da TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
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
Instalar o libtpu
pip install torch_xla[tpuvm] export XRT_TPU_CONFIG="localservice;0;localhost:51011"
Clonar o repositório do PyTorch XLA
git clone --recursive https://github.com/pytorch/xla.git
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.
- Digite
exit
para sair do contêiner do Docker - Digite
exit
para sair da VM da TPU - 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
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
SSH na VM da TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Iniciar o daemon do Docker na VM da TPU
sudo systemctl start docker
Iniciar contêiner do Docker
sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
Instalar o JAX
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Instalar o MXX
pip install --upgrade clu git clone https://github.com/google/flax.git pip install --user -e flax
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.
- Digite
exit
para sair do contêiner do Docker - Digite
exit
para sair da VM da TPU 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.
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árioyour-project-name
pelo nome do projeto do Google Cloudyour-repo-name
pelo nome do seu repositório do Artifact Registryyour-image-name
por um nome para sua imagem do Dockeryour-tag
por uma tag para a imagem do Dockeryour-tpu-name
por um nome para sua TPU
Crie um arquivo chamado
Dockerfile
no diretório atual e cole o texto a seguirFROM 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
Criar a imagem do Docker
docker build -t $IMAGE .
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
Enviar a imagem do Docker para o Artifact Registry
docker push europe-west4-docker.pkg.dev/ml-writers/$REPO_NAME/$IMAGE_NAME:$TAG
Criar uma VM de TPU
gcloud compute tpus tpu-vm create $TPU_NAME \ --zone $ZONE \ --accelerator-type $ACCEL_TYPE \ --version tpu-vm-base
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"
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"
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
- Tutoriais da Cloud TPU
- Guia do usuário do Cloud TPU
- Arquitetura do sistema do Cloud TPU
- Executar o código do TensorFlow em frações de Pod de TPU
- Executar código JAX em frações de Pod de TPU