Ejecuta aplicaciones de Cloud TPU en un contenedor de Docker.
Los contenedores de Docker configurar las aplicaciones con mayor facilidad combinando tu código y todos los elementos las dependencias en un paquete distribuible. Puedes ejecutar contenedores de Docker VM de TPU para simplificar la configuración y el uso compartido de tus aplicaciones de Cloud TPU. Esta que describe cómo configurar un contenedor de Docker para cada marco de trabajo de AA compatibles con Cloud TPU.
Entrena un modelo de TensorFlow en un contenedor de Docker
Dispositivo de TPU
Crea un archivo llamado
Dockerfile
en tu directorio actual y pega el texto siguienteFROM 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
Crea un bucket de Cloud Storage
gcloud storage buckets create gs://your-bucket-name --location=europe-west4
Crea una 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.17.0-pjrt
Copia el Dockerfile en tu VM de TPU
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
Establece una conexión SSH a la VM de TPU
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Compila la imagen de Docker
sudo docker build -t your-image-name .
Inicia el contenedor de Docker
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
Configura variables de entorno
export STORAGE_BUCKET=gs://your-bucket-name export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x
Entrena 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"
Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.
- Escribe
exit
para salir del contenedor de Docker. - Escribe
exit
para salir de la VM de TPU - Borra la VM de TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
pod de TPU
Crea un archivo llamado
Dockerfile
en tu directorio actual y pega el texto siguienteFROM 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
Crea una 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.17.0-pod-pjrt
Copia el Dockerfile en tu VM de TPU
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
Establece una conexión SSH a la VM de TPU
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Compila la imagen de Docker
sudo docker build -t your-image-name .
Inicia un contenedor de Docker
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
Entrena 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"
Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.
- Escribe
exit
para salir del contenedor de Docker. - Escribe
exit
para salir de la VM de TPU - Borra la VM de TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Entrena un modelo de PyTorch en un contenedor de Docker
Dispositivo de TPU
Crear una VM de Cloud TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-ubuntu2204-base
Establece una conexión SSH a la VM de TPU
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Inicia un contenedor en la VM de TPU con la imagen de PyTorch/XLA nocturna.
sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm bash
Configura el entorno de ejecución de la TPU
Existen dos opciones de entorno de ejecución de PyTorch/XLA: PJRT y XRT. Te recomendamos puede usar PJRT a menos que tenga un motivo. Para obtener más información sobre el diferentes configuraciones de entorno de ejecución, verás que tienes un motivo para usar XRT. Para Para obtener más información sobre las diferentes opciones de configuración del entorno de ejecución, consulta la documentación sobre el entorno de ejecución de PJRT.
PJRT
export PJRT_DEVICE=TPU
XRT
export XRT_TPU_CONFIG="localservice;0;localhost:51011"
Clona el repositorio de XLA de PyTorch
git clone --recursive https://github.com/pytorch/xla.git
Entrenar ResNet50
python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.
- Escribe
exit
para salir del contenedor de Docker. - Escribe
exit
para salir de la VM de TPU - Borra la VM de TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
pod de TPU
Cuando ejecutas código de PyTorch en un pod de TPU, debes hacerlo en todas las TPU
trabajadores al mismo tiempo. Una forma de hacerlo es con la
el comando gcloud compute tpus tpu-vm ssh
con --worker=all
y
--command
. En el siguiente procedimiento, se muestra cómo crear una instancia
para facilitar la configuración de cada trabajador TPU.
Crea una VM de TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=us-central2-b \ --accelerator-type=v4-32 \ --version=tpu-ubuntu2204-base
Agrega el usuario actual al grupo de Docker
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=us-central2-b \ --worker=all \ --command="sudo usermod -a -G docker $USER"
Ejecuta la secuencia de comandos de entrenamiento en un contenedor en todos los trabajadores 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"
Marcas de comando de Docker:
--rm
quita el contenedor después de que finaliza su proceso.--privileged
expone el dispositivo de TPU al contenedor.--net=host
vincula todos los puertos del contenedor a la VM de TPU para permitir la comunicación entre los hosts en el Pod.-e
configuró variables de entorno.
Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.
Borra la VM de TPU con el siguiente comando:
$ gcloud compute tpus tpu-vm delete your-tpu-name \
--zone=us-central2-b
Entrena un modelo de JAX en un contenedor de Docker
Dispositivo de TPU
Crea la VM de TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-ubuntu2204-base
Establece una conexión SSH a la VM de TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Iniciar el daemon de Docker en la VM de TPU
sudo systemctl start docker
Inicia el contenedor de Docker
sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
Instala 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
Ejecuta la secuencia de comandos de entrenamiento de 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
Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.
- Escribe
exit
para salir del contenedor de Docker. - Escribe
exit
para salir de la VM de TPU Borra la VM de TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
pod de TPU
Cuando ejecutas el código JAX en un Pod de TPU, debes ejecutar tu código JAX en todas las TPU
trabajadores al mismo tiempo. Una forma de hacerlo es usar gcloud compute tpus tpu-vm ssh
con las marcas --worker=all
y --command
. Lo siguiente
en el que se muestra cómo crear una imagen de Docker para configurar cada TPU
a los trabajadores de una manera más fácil.
Crea un archivo llamado
Dockerfile
en tu directorio actual y pega el texto siguienteFROM 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
Compila la imagen de Docker
docker build -t your-image-name .
Agrega una etiqueta a tu imagen de Docker antes de enviarla a Artifact Registry. Para obtener más información sobre cómo trabajar con Artifact Registry, consulta Trabaja con imágenes de contenedor.
docker tag your-image-name europe-west-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
Envía tu imagen de Docker a Artifact Registry
docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
Crea una VM de TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type==v2-8 \ --version=tpu-ubuntu2204-base
Extrae la imagen de Docker de Artifact Registry en todos los trabajadores 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"
Ejecuta el contenedor en todos los trabajadores de 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"
Ejecuta la secuencia de comandos de entrenamiento en todos los trabajadores 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"
Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.
Cierra el contenedor en todos los trabajadores:
gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \ --zone=europe-west4-a \ --command="docker kill your-container-name"
Borra la VM de TPU con el siguiente comando:
$ gcloud compute tpus tpu-vm delete your-tpu-name \ --zone=europe-west4-a
¿Qué sigue?
- Instructivos de Cloud TPU
- Administra las TPU
- Arquitectura de sistemas de Cloud TPU
- Ejecuta el código de TensorFlow en porciones de pod de TPU
- Ejecuta el código JAX en porciones de pod de TPU