Executar aplicativos do Cloud TPU em um contêiner do Docker
Os contêineres 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 da TPU para simplificar a configuração e o compartilhamento dos aplicativos da Cloud TPU. Neste documento, descrevemos como configurar um contêiner do Docker para cada framework de ML compatível com o 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 seguir:FROM python:3.8 RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.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 da TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-tf-2.13.0
Copie o Dockerfile para a 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
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"
Quando o script de treinamento for concluído, certifique-se de limpar 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
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.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.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 da TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v3-32 \ --version=tpu-vm-tf-2.13.0-pod
Copie o Dockerfile para a 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=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*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
Quando o script de treinamento for concluído, certifique-se de limpar 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
Dispositivo TPU
Criar VM do Cloud TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-ubuntu2204-base
SSH na VM da TPU
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Inicie um contêiner na VM da TPU usando a imagem noturna do PyTorch/XLA.
sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm bash
Configurar o ambiente de execução da TPU
Há duas opções de ambiente de execução do PyTorch/XLA: PJRT e XRT. Recomendamos que você use o PJRT, a menos que tenha um motivo. Para saber mais sobre as diferentes configurações de ambiente de execução, veja por que você precisa usar o XRT. Para saber mais sobre as diferentes configurações do ambiente de execução, consulte a documentação do ambiente de execução do PJRT.
PJRT
export PJRT_DEVICE=TPU
XRT
export XRT_TPU_CONFIG="localservice;0;localhost:51011"
Clone o repositório PyTorch XLA
git clone --recursive https://github.com/pytorch/xla.git
Treinar ResNet50
python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
Quando o script de treinamento for concluído, certifique-se de limpar 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 PyTorch em um Pod de TPU, é preciso executar o 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 da TPU.
Criar uma VM da TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=us-central2-b \ --accelerator-type=v4-32 \ --version=tpu-ubuntu2204-base
Adicionar o usuário atual ao grupo do Docker
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=us-central2-b \ --worker=all \ --command="sudo usermod -a -G docker $USER"
Execute o script de treinamento em um contêiner em todos os workers da TPU.
gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \ --zone=us-central2-b \ --command="docker run --rm --privileged --net=host -e PJRT_DEVICE=TPU gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm python /pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"
Sinalizações de comando do Docker:
--rm
remove o contêiner após o término do processo.--privileged
expõe o dispositivo da TPU ao contêiner.--net=host
vincula todas as portas do contêiner à VM da TPU para permitir a comunicação entre os hosts no pod.-e
definiu variáveis de ambiente.
Quando o script de treinamento for concluído, certifique-se de limpar os recursos.
Exclua a VM da TPU usando o seguinte comando:
$ gcloud compute tpus tpu-vm delete your-tpu-name \
--zone=us-central2-b
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-ubuntu2204-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] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Instalar FLAX
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
Quando o script de treinamento for concluído, certifique-se de limpar 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, é necessário executá-lo 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 da TPU.
Crie um arquivo chamado
Dockerfile
no diretório atual e cole o texto a seguir:FROM python:3.8 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 your-image-name .
Adicione uma tag à imagem do Docker antes de enviá-la ao Artifact Registry. Para mais informações sobre como trabalhar com o Artifact Registry, consulte Trabalhar com imagens de contêiner.
docker tag your-image-name europe-west-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
Enviar a imagem do Docker para o Artifact Registry
docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
Criar uma VM da TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type==v2-8 \ --version=tpu-ubuntu2204-base
Extraia a imagem do Docker do Artifact Registry em todos os workers da TPU.
gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \ --zone=europe-west4-a \ --command="sudo usermod -a -G docker ${USER}"
gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \ --zone=europe-west4-a \ --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \ --zone=europe-west4-a \ --command="docker pull europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag"
Execute o contêiner em todos os workers da TPU.
gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \ zone=europe-west4-a \ --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image:your-tag bash"
Execute o script de treinamento em todos os workers da TPU:
gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \ --zone=europe-west4-a \ --command="docker exec --privileged your-container-name python3 main.py --workdir=/tmp/mnist \ --config=configs/default.py \ --config.learning_rate=0.05 \ --config.num_epochs=5"
Quando o script de treinamento for concluído, certifique-se de limpar os recursos.
Encerre o contêiner em todos os workers:
gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \ --zone=europe-west4-a \ --command="docker kill your-container-name"
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
- Gerenciar TPUs
- Arquitetura do sistema do Cloud TPU
- Executar o código do TensorFlow em Pod de TPU TPU
- Executar o código JAX em Pod de TPU TPU