Entrenamiento de Cloud TPU v5e

Cloud TPU v5e es el acelerador de IA de nueva generación de Google Cloud. Con una huella más pequeña de 256 chips por pod, un v5e está optimizado para ser el producto de mayor valor para el entrenamiento, el perfeccionamiento y la publicación de transformadores, texto a imagen y redes neuronales convolucionales (CNN). Para obtener más información sobre el uso de Cloud TPU v5e para la publicación, consulta Inferencia con v5e.

Para obtener más información sobre el hardware y las configuraciones de TPU de Cloud TPU v5e, consulta TPU v5e.

Comenzar

En las siguientes secciones, se describe cómo comenzar a usar las TPU v5e.

Cuota de solicitudes

Necesitas una cuota para usar TPU v5e para el entrenamiento. Existen diferentes tipos de cuotas para las TPUs bajo demanda, las TPU reservadas y las VMs de TPU Spot. Se requieren cuotas separadas si usas tu TPU v5e para la inferencia. Para obtener más información sobre las cuotas, consulta Cuotas. Para solicitar una cuota de TPU v5e, comunícate con Ventas de Cloud.

Crea una cuenta y un proyecto de Google Cloud

Necesitas una cuenta y un proyecto de Google Cloud para usar Cloud TPU. Para obtener más información, consulta Configura un entorno de Cloud TPU.

Crear una Cloud TPU

La práctica recomendada es aprovisionar Cloud TPU v5 como recursos en fila con el comando queued-resource create. Para obtener más información, consulta Administra recursos en fila.

También puedes usar la API de Create Node (gcloud compute tpus tpu-vm create) para aprovisionar Cloud TPU v5. Para obtener más información, consulta Administra recursos de TPU.

Para obtener más información sobre las configuraciones v5e disponibles para el entrenamiento, consulta Tipos de Cloud TPU v5e para el entrenamiento.

Configuración del framework

En esta sección, se describe el proceso de configuración general para el entrenamiento de modelos personalizados con JAX o PyTorch con TPU v5e. La compatibilidad con TensorFlow está disponible en las versiones del entorno de ejecución de TPU tpu-vm-tf-2.18.0-pjrt y tpu-vm-tf-2.18.0-pod-pjrt.

Para obtener instrucciones de configuración de la inferencia, consulta la introducción a la inferencia de v5e.

Configuración para JAX

Si tienes formas de rebanada superiores a 8 chips, tendrás varias VMs en una rebanada. En este caso, debes usar la marca --worker=all para ejecutar la instalación en todas las VMs de TPU en un solo paso sin usar SSH para acceder a cada una por separado:

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Descripciones de las marcas de comandos

Variable Descripción
TPU_NAME El ID de texto asignado por el usuario de la TPU que se crea cuando se asigna la solicitud de recurso en cola.
ID DEL PROYECTO Nombre del proyecto de Google Cloud. Usa un proyecto existente o crea uno nuevo en Configura tu proyecto de Google Cloud.
ZONA Consulta el documento Regiones y zonas de TPU para conocer las zonas compatibles.
Worker La VM de TPU que tiene acceso a las TPU subyacentes.

Puedes ejecutar el siguiente comando para verificar la cantidad de dispositivos (los resultados que se muestran aquí se produjeron con una porción v5litepod-16). Este código prueba que todo está instalado correctamente, ya que verifica que JAX vea los TensorCores de Cloud TPU y pueda ejecutar operaciones básicas:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

El resultado será similar al siguiente ejemplo:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() muestra la cantidad total de chips en la porción determinada. jax.local_device_count() indica la cantidad de chips a los que puede acceder una sola VM en esta porción.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

El resultado será similar al siguiente ejemplo:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Prueba los instructivos de JAX en este documento para comenzar a usar el entrenamiento de v5e con JAX.

Configuración para PyTorch

Ten en cuenta que la v5e solo admite el entorno de ejecución de PJRT, y PyTorch 2.1 y versiones posteriores usarán PJRT como el entorno de ejecución predeterminado para todas las versiones de TPU.

En esta sección, se describe cómo comenzar a usar PJRT en v5e con PyTorch/XLA con comandos para todos los trabajadores.

Instala dependencias

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

Si recibes un error cuando instalas las ruedas para torch, torch_xla o torchvision como pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, revierte la versión con este comando:

pip3 install setuptools==62.1.0

Ejecuta una secuencia de comandos con PJRT

unset LD_PRELOAD

A continuación, se muestra un ejemplo en el que se usa una secuencia de comandos de Python para realizar un cálculo en una VM v5e:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

Esto genera un resultado similar al que se muestra a continuación:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

Prueba los instructivos de PyTorch que se incluyen en este documento para comenzar a usar el entrenamiento de v5e con PyTorch.

Borra tu TPU y el recurso en cola al final de la sesión. Para borrar un recurso en cola, borra la porción y, luego, el recurso en cola en 2 pasos:

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Estos dos pasos también se pueden usar para quitar las solicitudes de recursos en fila que están en el estado FAILED.

Ejemplos de JAX/FLAX

En las siguientes secciones, se describen ejemplos de cómo entrenar modelos de JAX y FLAX en TPU v5e.

Entrena ImageNet en v5e

En este instructivo, se describe cómo entrenar ImageNet en la versión 5e con datos de entrada falsos. Si quieres usar datos reales, consulta el archivo README en GitHub.

Configurar

  1. Crea variables de entorno:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Crea un recurso de TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Podrás establecer una conexión SSH a tu VM de TPU una vez que el recurso en cola esté en el estado ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Cuando QueuedResource esté en el estado ACTIVE, el resultado será similar al siguiente:

     state: ACTIVE

  3. Instala la versión más reciente de JAX y jaxlib:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Clona el modelo de ImageNet y, luego, instala los requisitos correspondientes:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'

  5. Para generar datos falsos, el modelo necesita información sobre las dimensiones del conjunto de datos. Esto se puede obtener de los metadatos del conjunto de datos de ImageNet:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'

Entrena el modelo

Una vez que hayas completado todos los pasos anteriores, podrás entrenar el modelo.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

Borra la TPU y el recurso en cola

Borra tu TPU y el recurso en cola al final de la sesión.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Modelos FLAX de Hugging Face

Los modelos de Hugging Face implementados en FLAX funcionan de forma predeterminada en Cloud TPU v5e. En esta sección, se proporcionan instrucciones para ejecutar modelos populares.

Entrena ViT en Imagenette

En este instructivo, se muestra cómo entrenar el modelo Vision Transformer (ViT) de HuggingFace con el conjunto de datos Imagenette de Fast AI en Cloud TPU v5e.

El modelo de ViT fue el primero que entrenó con éxito un codificador de Transformer en ImageNet con resultados excelentes en comparación con las redes convolucionales. Para obtener más información, consulta los siguientes recursos:

Configurar

  1. Crea variables de entorno:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Crea un recurso de TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Podrás establecer una conexión SSH a tu VM de TPU una vez que el recurso en cola esté en el estado ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Cuando el recurso en cola esté en el estado ACTIVE, el resultado será similar al siguiente:

     state: ACTIVE

  3. Instala JAX y su biblioteca:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Descarga el repositorio de Hugging Face y los requisitos de instalación:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.18.0 && pip install -r examples/flax/vision/requirements.txt'

  5. Descarga el conjunto de datos de Imagenette:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'

Entrena el modelo

Entrena el modelo con un búfer preasignado de 4 GB.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

Borra la TPU y el recurso en cola

Borra la TPU y el recurso en cola al final de la sesión.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Resultados de las comparativas de ViT

La secuencia de comandos de entrenamiento se ejecutó en v5litepod-4, v5litepod-16 y v5litepod-64. En la siguiente tabla, se muestran las tasas de transferencia con diferentes tipos de aceleradores.

Tipo de acelerador v5litepod-4 v5litepod-16 v5litepod-64
Epoch 3 3 3
Tamaño del lote global 32 128 512
Capacidad de procesamiento (ejemplos/s) 263.40 429.34 470.71

Entrena la difusión en Pokémon

En este instructivo, se muestra cómo entrenar el modelo de difusión estable de HuggingFace con el conjunto de datos de Pokémon en Cloud TPU v5e.

El modelo de difusión estable es un modelo latente de texto a imagen que genera imágenes fotorrealistas a partir de cualquier entrada de texto. Para obtener más información, consulta los siguientes recursos:

Configurar

  1. Configura un bucket de almacenamiento para el resultado de tu modelo. 

    gcloud storage buckets create gs://your_bucket 
    
--project=your_project 
    
--location=us-west1 
    
export GCS_BUCKET_NAME=your_bucket

  2. Crea variables de entorno

    export GCS_BUCKET_NAME=your_bucket
export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west1-c
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

    Descripciones de las marcas de comandos

    Variable Descripción
    GCS_BUCKET_NAME Se muestra en la consola de Google Cloud -> Cloud Storage -> Buckets
    ID DEL PROYECTO Nombre del proyecto de Google Cloud. Usa un proyecto existente o crea uno nuevo en Configura tu proyecto de Google Cloud.
    ACCELERATOR_TYPE Consulta la página Versiones de TPU para ver la versión de TPU.
    ZONA Consulta el documento Regiones y zonas de TPU para conocer las zonas compatibles.
    RUNTIME_VERSION Usa v2-alpha-tpuv5 para RUNTIME_VERSION.
    SERVICE_ACCOUNT Esta es la dirección de tu cuenta de servicio que puedes encontrar en la consola de Google Cloud -> IAM -> Cuentas de servicio. Por ejemplo: tpu-service-account@myprojectID.iam.gserviceaccount.com
    TPU_NAME El ID de texto asignado por el usuario de la TPU que se crea cuando se asigna la solicitud de recursos en cola.
    QUEUED_RESOURCE_ID El ID de texto asignado por el usuario de la solicitud de recursos en cola. Consulta el documento Recursos en cola para obtener información sobre los recursos en cola.
    QUOTA_TYPE Puede ser reserved o spot. Si no se especifica ninguno de estos, el valor predeterminado de QUOTA_TYPE es on-demand. Consulta cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.
    VALID_UNTIL_DURATION Es la duración de la validez de la solicitud. Consulta Recursos en fila para obtener información sobre las diferentes duraciones válidas.

  3. Crea un recurso de TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Podrás establecer una conexión SSH a tu VM de TPU una vez que el recurso en cola esté en el estado ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Cuando el recurso en cola esté en el estado ACTIVE, el resultado será similar al siguiente:

     state: ACTIVE

  4. Instala JAX y su biblioteca.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  5. Descarga el repositorio de HuggingFace y los requisitos de instalación.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.18.0 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'

Entrena el modelo

Entrena el modelo con un búfer preasignado de 4 GB.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE} --project ${PROJECT_ID} --worker=all --command="
   git clone https://github.com/google/maxdiffusion
   cd maxdiffusion
   git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
   pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   pip3 install -r requirements.txt
   pip3 install .
   export LIBTPU_INIT_ARGS=""
   python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=your_run base_output_directory=gs://${GCS_BUCKET_NAME}/ enable_profiler=False"

Borra la TPU y el recurso en cola

Borra tu TPU y el recurso en cola al final de la sesión.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Resultados de las comparativas de difusión

La secuencia de comandos de entrenamiento se ejecutó en v5litepod-4, v5litepod-16 y v5litepod-64. En la siguiente tabla, se muestran las capacidades de procesamiento.

Tipo de acelerador v5litepod-4 v5litepod-16 v5litepod-64
Paso de entrenamiento 1500 1500 1500
Tamaño del lote global 32 64 128
Capacidad de procesamiento (ejemplos/s) 36.53 43.71 49.36

Entrena GPT2 en el conjunto de datos de OSCAR

En este instructivo, se muestra cómo entrenar el modelo GPT2 de HuggingFace con el conjunto de datos OSCAR en Cloud TPU v5e.

GPT2 es un modelo de transformador previamente entrenado en textos sin procesar sin etiquetado humano. Se entrenó para predecir la siguiente palabra en oraciones. Para obtener más información, consulta los siguientes recursos:

Configurar

  1. Crea variables de entorno:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Crea un recurso de TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Podrás establecer una conexión SSH a tu VM de TPU una vez que el recurso en cola esté en el estado ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Cuando el recurso en cola esté en el estado ACTIVE, el resultado será similar al siguiente:

     state: ACTIVE

  3. Instala JAX y su biblioteca.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Descarga el repositorio de HuggingFace y, luego, instala los requisitos.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install TensorFlow && pip install -r examples/flax/language-modeling/requirements.txt'

  5. Descarga las configuraciones para entrenar el modelo.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt . --recursive'

Entrena el modelo

Entrena el modelo con un búfer preasignado de 4 GB.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'

Borra la TPU y el recurso en cola

Borra tu TPU y el recurso en cola al final de la sesión.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Resultados de las comparativas de GPT2

La secuencia de comandos de entrenamiento se ejecutó en v5litepod-4, v5litepod-16 y v5litepod-64. En la siguiente tabla, se muestran las capacidades de procesamiento.

v5litepod-4 v5litepod-16 v5litepod-64
Epoch 3 3 3
Tamaño del lote global 64 64 64
Capacidad de procesamiento (ejemplos/s) 74.60 72.97 72.62

PyTorch/XLA

En las siguientes secciones, se describen ejemplos de cómo entrenar modelos de PyTorch/XLA en TPU v5e.

Entrena ResNet con el entorno de ejecución de PJRT

PyTorch/XLA migra de XRT a PjRt desde PyTorch 2.0 y versiones posteriores. Aquí tienes las instrucciones actualizadas para configurar v5e para las cargas de trabajo de entrenamiento de PyTorch/XLA.

Configurar

  1. Crea variables de entorno:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Crea un recurso de TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --{QUOTA_TYPE}

    Podrás establecer una conexión SSH a tu VM de TPU una vez que QueuedResource esté en el estado ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Cuando el recurso en cola esté en el estado ACTIVE, el resultado será similar al siguiente:

     state: ACTIVE

  3. Instala dependencias específicas de Torch/XLA

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Entrena el modelo ResNet
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

Borra la TPU y el recurso en cola

Borra tu TPU y el recurso en cola al final de la sesión.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet
Resultado de la comparativa

En la siguiente tabla, se muestran las capacidades de procesamiento de comparativas.

Tipo de acelerador Capacidad de procesamiento (ejemplos por segundo)
v5litepod-4 4240 ex/s
v5litepod-16 10,810 ex/s
v5litepod-64 46,154 ex/s

Entrena GPT2 en v5e

En este instructivo, se explica cómo ejecutar GPT2 en v5e con el repositorio de HuggingFace en PyTorch/XLA con el conjunto de datos de wikitext.

Configurar

  1. Crea variables de entorno:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Crea un recurso de TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Podrás establecer una conexión SSH a tu VM de TPU una vez que tu elemento QueuedResource esté en estado ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Cuando el recurso en cola esté en el estado ACTIVE, el resultado será similar al siguiente:

    state: ACTIVE

  3. Instala las dependencias de PyTorch/XLA.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get -y update
      sudo apt install -y libopenblas-base
      pip3 install torchvision
      pip3 uninstall -y torch
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

  4. Descarga el repositorio de HuggingFace y, luego, instala los requisitos.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      git clone https://github.com/pytorch/xla.git
      pip install --upgrade accelerate
      git clone https://github.com/huggingface/transformers.git
      cd transformers
      git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
      pip install .
      git log -1
      pip install datasets evaluate scikit-learn
      '

  5. Descargar las configuraciones del modelo previamente entrenado

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/ --recursive
      gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'

Entrena el modelo

Entrena el modelo 2B con un tamaño de lote de 16.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 examples/pytorch/xla_spawn.py \
         --num_cores=4 \
         examples/pytorch/language-modeling/run_clm.py \
         --num_train_epochs=3 \
         --dataset_name=wikitext \
         --dataset_config_name=wikitext-2-raw-v1 \
         --per_device_train_batch_size=16 \
         --per_device_eval_batch_size=16 \
         --do_train \
         --do_eval \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --tokenizer_name=gpt2 \
         --block_size=1024 \
         --optim=adafactor \
         --adafactor=true \
         --save_strategy=no \
         --logging_strategy=no \
         --fsdp=full_shard \
         --fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'

Borra la TPU y el recurso en cola

Borra tu TPU y el recurso en cola al final de la sesión.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Resultado de la comparativa

La secuencia de comandos de entrenamiento se ejecutó en v5litepod-4, v5litepod-16 y v5litepod-64. En la siguiente tabla, se muestran las tasas de rendimiento de las comparativas para diferentes tipos de aceleradores.

v5litepod-4 v5litepod-16 v5litepod-64
Epoch 3 3 3
config 600 millones 2,000 MILLONES 16,000 millones
Tamaño del lote global 64 128 256
Capacidad de procesamiento (ejemplos/s) 66 77 31

Entrena ViT en v5e

En este instructivo, se explica cómo ejecutar VIT en v5e con el repositorio de HuggingFace en PyTorch/XLA en el conjunto de datos cifar10.

Configurar

  1. Crea variables de entorno:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Crea un recurso de TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Podrás establecer una conexión SSH a tu VM de TPU una vez que QueuedResource esté en el estado ACTIVE:

     gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Cuando el recurso en cola esté en el estado ACTIVE, el resultado será similar al siguiente:

     state: ACTIVE

  3. Instala las dependencias de PyTorch/XLA

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

  4. Descarga el repositorio de HuggingFace y, luego, instala los requisitos.

       gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="
      git clone https://github.com/suexu1025/transformers.git vittransformers; \
      cd vittransformers; \
      pip3 install .; \
      pip3 install datasets; \
      wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"

Entrena el modelo

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
      --remove_unused_columns=False \
      --label_names=pixel_values \
      --mask_ratio=0.75 \
      --norm_pix_loss=True \
      --do_train=true \
      --do_eval=true \
      --base_learning_rate=1.5e-4 \
      --lr_scheduler_type=cosine \
      --weight_decay=0.05 \
      --num_train_epochs=3 \
      --warmup_ratio=0.05 \
      --per_device_train_batch_size=8 \
      --per_device_eval_batch_size=8 \
      --logging_strategy=steps \
      --logging_steps=30 \
      --evaluation_strategy=epoch \
      --save_strategy=epoch \
      --load_best_model_at_end=True \
      --save_total_limit=3 \
      --seed=1337 \
      --output_dir=MAE \
      --overwrite_output_dir=true \
      --logging_dir=./tensorboard-metrics \
      --tpu_metrics_debug=true'

Borra la TPU y el recurso en cola

Borra tu TPU y el recurso en cola al final de la sesión.

gcloud compute tpus tpu-vm delete ${TPU_NAME}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

Resultado de la comparativa

En la siguiente tabla, se muestran las tasas de rendimiento de las comparativas para diferentes tipos de aceleradores.

v5litepod-4 v5litepod-16 v5litepod-64
Epoch 3 3 3
Tamaño del lote global 32 128 512
Capacidad de procesamiento (ejemplos/s) 201 657 2,844

TensorFlow 2.x

En las siguientes secciones, se describen ejemplos de cómo entrenar modelos de TensorFlow 2.x en TPU v5e.

Entrena Resnet en una v5e de host único

En este instructivo, se describe cómo entrenar ImageNet en v5litepod-4 o v5litepod-8 con un conjunto de datos falso. Si quieres usar un conjunto de datos diferente, consulta Cómo preparar el conjunto de datos.

Configurar

  1. Crea variables de entorno:

    export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5litepod-4
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE puede ser v5litepod-4 o v5litepod-8.

  2. Crea un recurso de TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --${QUOTA_TYPE}

    Podrás establecer una conexión SSH a tu VM de TPU una vez que el recurso en cola esté en el estado ACTIVE. Para verificar el estado de tu recurso en cola, usa el siguiente comando:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. Conéctate a tu TPU con SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  4. Configura algunas variables de entorno

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

  5. Cambia al directorio del repositorio de modelos y, luego, instala los requisitos.

    cd ${MODELS_REPO} && git checkout r2.15.0
pip install -r official/requirements.txt

Entrena el modelo

Ejecuta la secuencia de comandos de entrenamiento.

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="runtime.distribution_strategy=tpu,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"

Borra la TPU y el recurso en cola

  1. Borra tu TPU

    gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

  2. Cómo borrar tu solicitud de recursos en cola

    gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Entrena Resnet en un v5e multihost

En este instructivo, se describe cómo entrenar ImageNet en v5litepod-16 o más con un conjunto de datos falso. Si quieres usar un conjunto de datos diferente, consulta Prepara el conjunto de datos.

  1. Crea variables de entorno:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE puede ser v5litepod-16 o mayor.

  2. Crea un recurso de TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --${QUOTA_TYPE}

    Podrás establecer una conexión SSH a tu VM de TPU una vez que el recurso en cola esté en el estado ACTIVE. Para verificar el estado de tu recurso en cola, usa el siguiente comando:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. Conéctate a tu TPU (trabajador cero) con SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  4. Configura algunas variables de entorno

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export TPU_LOAD_LIBRARY=0
export TPU_NAME=your_tpu_name

  5. Cambia al directorio del repositorio de modelos y, luego, instala los requisitos.

     cd $MODELS_REPO && git checkout r2.15.0
 pip install -r official/requirements.txt

Entrena el modelo

Ejecuta la secuencia de comandos de entrenamiento.

python3 official/vision/train.py \
   --tpu=${TPU_NAME} \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*"

Borra la TPU y el recurso en cola

  1. Borra tu TPU

    gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

  2. Cómo borrar tu solicitud de recursos en cola

    gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet