Entrena a un transformador en Cloud TPU (TF 2.x)


Si no conoces Cloud TPU, consulta las guías de introducción a TPU.

En este instructivo, se muestra cómo entrenar un modelo Transformer en Cloud TPU. Transformer es una arquitectura de red neuronal que resuelve problemas de secuencia a secuencia con mecanismos de atención. A diferencia de los modelos tradicionales de seq2seq, El transformador no implica conexiones recurrentes. El mecanismo de atención aprende dependencias entre tokens en dos secuencias. Dado que las ponderaciones de atención se aplican a todos los tokens de las secuencias, el modelo de Transformer puede capturar fácilmente las dependencias de larga distancia.

La estructura general de Transformer sigue el patrón de codificador-decodificador estándar. El codificador usa la autoatención para procesar una representación de la entrada secuencia. El decodificador genera la secuencia de salida un token a la vez y toma la salida del codificador y los tokens de salida del decodificador anteriores como entradas.

El modelo también aplica incorporaciones a los tokens de entrada y salida, y agrega codificación posicional constante. La codificación posicional agrega información sobre la posición de cada token.

Objetivos

  • Crear un bucket de Cloud Storage para almacenar el resultado de tu modelo y tu conjunto de datos
  • Descarga y procesa previamente el conjunto de datos que se usa para entrenar el modelo.
  • Ejecutar el trabajo de entrenamiento
  • Verificar los resultados de salida

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

  • Compute Engine
  • Cloud TPU

Para generar una estimación de costos en función del uso previsto, usa la calculadora de precios. Es posible que los usuarios nuevos de Google Cloud califiquen para obtener una prueba gratuita.

Antes de comenzar

Antes de comenzar este instructivo, verifica que tu proyecto de Google Cloud esté configurado correctamente.

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

  6. En esta explicación, se usan componentes facturables de Google Cloud. Consulta la página de precios de Cloud TPU para calcular los costos. Asegúrate de limpiar los recursos que crees cuando hayas terminado de usarlos para evitar cargos innecesarios.

Entrena con un solo dispositivo de Cloud TPU

Esta sección brinda información sobre cómo configurar un bucket de Cloud Storage y una VM de TPU para el entrenamiento de un solo dispositivo.

  1. Abre una ventana de Cloud Shell.

    Abra Cloud Shell

  2. Crea una variable de entorno para el ID de tu proyecto.

    export PROJECT_ID=project-id
  3. Configura Google Cloud CLI para usar tu proyecto de Google Cloud en el que deseas crear una Cloud TPU.

    gcloud config set project ${PROJECT_ID}

    La primera vez que ejecutas este comando en una VM nueva de Cloud Shell, se muestra una página Authorize Cloud Shell. Haz clic en Authorize en la parte inferior de la página para permitir que gcloud realice llamadas a la API de Google Cloud con tus credenciales.

  4. Crea una cuenta de servicio para el proyecto de Cloud TPU.

    Las cuentas de servicio permiten que el servicio de Cloud TPU acceda a otros Google Cloud de Google Cloud.

    $ gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID

    El comando muestra una cuenta de servicio de Cloud TPU con el siguiente formato:

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

  5. Crea un bucket de Cloud Storage con el siguiente comando:

    $ gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=us-central2

    En este bucket de Cloud Storage se almacenan los datos que usas para entrenar tu modelo y los resultados del entrenamiento. Con el comando gcloud que se usa en este instructivo para configurar la TPU, también se configuran los permisos predeterminados para la cuenta de servicio de Cloud TPU que configuraste en el paso anterior. Si quieres contar con permisos más detallados, revisa los permisos de nivel de acceso.

Entrena el modelo de Transformer en una sola Cloud TPU

  1. Inicia una VM de Cloud TPU con el comando gcloud.

      $ gcloud compute tpus tpu-vm create transformer-tutorial \
         --zone=us-central2-b \
         --accelerator-type=v4-8 \
         --version=tpu-vm-tf-2.17.0-pjrt

    Descripciones de las marcas de comandos

    zone
    Es la zona en la que deseas crear la Cloud TPU.
    accelerator-type
    El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
    version
    La versión de software de Cloud TPU.
  2. Conéctate a la VM de TPU mediante SSH. Cuando estés conectado a la VM, el símbolo del sistema de shell cambia de username@projectname a username@vm-name:

    gcloud compute tpus tpu-vm ssh transformer-tutorial --zone=us-central2-b
  3. Exporta variables de entorno.

      (vm)$ export STORAGE_BUCKET=gs://bucket-name
      (vm)$ export SENTENCEPIECE_MODEL=sentencepiece
      (vm)$ export SENTENCEPIECE_MODEL_PATH=${STORAGE_BUCKET}/${SENTENCEPIECE_MODEL}.model
      (vm)$ export TFDS_DIR=${STORAGE_BUCKET}/tfds
      (vm)$ export PARAM_SET=big
      (vm)$ export TPU_NAME=local
      (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}
      (vm)$  export PYTHONPATH="/usr/share/tpu/models:$PYTHONPATH"
  4. Cuando crees tu TPU, si configuras el parámetro --version en una versión que finaliza con -pjrt, establece las siguientes variables de entorno para habilitar el entorno de ejecución de PJRT:

      (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
      (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
  5. Instala los requisitos de TensorFlow.

    (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt
  6. Descarga y procesa previamente el conjunto de datos

     (vm)$ python3 -c "import tensorflow_datasets as tfds; tfds.load('wmt14_translate/de-en', split='train+validation', shuffle_files=True, download=True)"
      (vm)$ python3 /usr/share/tpu/models/official/nlp/data/train_sentencepiece.py --output_model_path=${SENTENCEPIECE_MODEL}
  7. Copia el conjunto de datos en el bucket de Cloud Storage

      (vm)$ gcloud storage cp ${SENTENCEPIECE_MODEL}.model ${STORAGE_BUCKET}
      (vm)$ gcloud storage cp ${SENTENCEPIECE_MODEL}.vocab ${STORAGE_BUCKET}
      (vm)$ gcloud storage cp tensorflow_datasets/wmt14_translate ${TFDS_DIR}/wmt14_translate --recursive
  8. Navega al directorio de entrenamiento

    (vm)$ cd /usr/share/tpu/models/
  9. Ejecuta la secuencia de comandos de entrenamiento

    (vm)$ python3 official/nlp/train.py \
      --tpu=${TPU_NAME} \
      --experiment=wmt_transformer/large \
      --mode=train_and_eval \
      --model_dir=${MODEL_DIR} \
      --params_override="runtime.distribution_strategy=tpu, task.train_data.tfds_data_dir=${TFDS_DIR}, task.validation_data.tfds_data_dir=${TFDS_DIR}, task.sentencepiece_model_path=${SENTENCEPIECE_MODEL_PATH}, trainer.train_steps=10000, trainer.validation_interval=10000"

    Descripciones de las marcas de comandos

    tpu
    El nombre de Cloud TPU. Para ello, especifica la variable de entorno (TPU_NAME).
    experiment
    El modelo que se entrenará.
    mode
    Es el modo en el que se ejecutará la secuencia de comandos.
    model_dir
    El bucket de Cloud Storage en el que se almacenan los puntos de control y los resúmenes durante el entrenamiento. Puedes usar una carpeta existente para cargar puntos de control generados con anterioridad creados en una TPU del mismo tamaño y con la misma versión de TensorFlow.
    params_override
    Establece los parámetros del modelo.

    De forma predeterminada, el modelo se evaluará después de cada 10,000 pasos. Puedes aumentar la cantidad de pasos de entrenamiento o especificar con qué frecuencia ejecutar evaluaciones configurando estos parámetros:

    • train.train_steps: Es la cantidad total de pasos de entrenamiento que se deben ejecutar.
    • trainer.validation_interval: Es la cantidad de pasos de entrenamiento que se deben ejecutar entre las evaluaciones.

    El entrenamiento y la evaluación toman aproximadamente 20 minutos en una Cloud TPU v4-8. Cuando se completan el entrenamiento y la evaluación, aparece un mensaje similar al siguiente:

    I0208 20:57:19.309512 140039467895872 controller.py:310]  eval | step:  10000 | eval time:   69.2 sec | output:
    {'bleu_score': 19.204771518707275,
     'sacrebleu_score': 18.307039308307356,
     'validation_loss': 2.0654342}
     eval | step:  10000 | eval time:   69.2 sec | output:
    {'bleu_score': 19.204771518707275,
     'sacrebleu_score': 18.307039308307356,
     'validation_loss': 2.0654342}
    

    Ya completaste el entrenamiento con un solo dispositivo. Sigue estos pasos para y borrar los recursos de TPU de un solo dispositivo.

  10. Desconéctate de la instancia de Compute Engine:

    (vm)$ exit

    El mensaje ahora debería mostrar username@projectname, que indica que estás en Cloud Shell.

  11. Borra el recurso TPU.

    $ gcloud compute tpus tpu-vm delete transformer-tutorial \
    --zone=us-central2-b

    Descripciones de las marcas de comandos

    zone
    Es la zona en la que reside tu Cloud TPU.

En este punto, puedes finalizar este instructivo y realizar una limpieza o puedes continuar y explorar la ejecución del modelo en los Pods de Cloud TPU.

Escala tu modelo con pods de Cloud TPU

Entrenar tu modelo en pods de Cloud TPU puede requerir algunos cambios en el guion de entrenamiento. Para obtener más información, consulta Entrenamiento en pods de TPU.

Entrenamiento de Pods de TPU

  1. Abre una ventana de Cloud Shell.

    Abra Cloud Shell

  2. Crea una variable para el ID de tu proyecto.

    export PROJECT_ID=project-id
  3. Configura Google Cloud CLI para usar el proyecto en el que deseas crear una Cloud TPU.

    gcloud config set project ${PROJECT_ID}

    La primera vez que ejecutas este comando en una VM nueva de Cloud Shell, se muestra una página Authorize Cloud Shell. Haz clic en Authorize en la parte inferior de la página para permitir que gcloud realice llamadas a la API con tus credenciales.

  4. Crea una cuenta de servicio para el proyecto de Cloud TPU.

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID

    El comando muestra una cuenta de servicio de Cloud TPU con el siguiente formato:

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

  5. Crea un bucket de Cloud Storage mediante el siguiente comando o usa un bucket que hayas creado antes para tu proyecto:

    gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=us-central1

    En este bucket de Cloud Storage se almacenan los datos que usas para entrenar tu modelo y los resultados del entrenamiento. Con el comando gcloud que se usa en este instructivo, se configuran los permisos predeterminados para la cuenta de servicio de Cloud TPU que configuraste en el paso anterior. Para obtener más información sobre permisos más detallados, consulta permisos de nivel de acceso.

Inicia los recursos de la VM de TPU

  1. Inicia un pod de VM de TPU con el comando gcloud. En este instructivo, se especifica un Pod v4-32. Para conocer otras opciones de Pod, consulta la página de tipos de TPU disponibles de tipos de TPU.

    $ gcloud compute tpus tpu-vm create transformer-tutorial \
      --zone=us-central2-b \
      --accelerator-type=v3-32 \
      --version=tpu-vm-tf-2.17.0-pod-pjrt

    Descripciones de las marcas de comandos

    zone
    Es la zona en la que deseas crear la Cloud TPU.
    accelerator-type
    El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
    version
    La versión de software de Cloud TPU.
  2. Conéctate a la instancia de Compute Engine con SSH. Cuando estés conectado a la VM, el símbolo del sistema de shell cambia de username@projectname a username@vm-name:

    gcloud compute tpus tpu-vm ssh transformer-tutorial --zone=us-central2-b
  3. Instala los requisitos de TensorFlow.

    (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt

Configura y comienza el entrenamiento del pod

  1. Exporta las variables de configuración de Cloud TPU:

    (vm)$ export PYTHONPATH="/usr/share/tpu/models:$PYTHONPATH"
    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    (vm)$ export SENTENCEPIECE_MODEL=sentencepiece
    (vm)$ export SENTENCEPIECE_MODEL_PATH=${STORAGE_BUCKET}/${SENTENCEPIECE_MODEL}.model
    (vm)$ export TFDS_DIR=${STORAGE_BUCKET}/tfds
    (vm)$ export TPU_NAME=transformer-tutorial
    (vm)$ export PARAM_SET=big
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}
    (vm)$ export TPU_LOAD_LIBRARY=0
  2. Descarga el conjunto de datos

    (vm)$ python3 -c "import tensorflow_datasets as tfds; tfds.load('wmt14_translate/de-en', split='train+validation', shuffle_files=True, download=True)"
    (vm)$ python3 /usr/share/tpu/models/official/nlp/data/train_sentencepiece.py --output_model_path=${SENTENCEPIECE_MODEL}
  3. Copia el conjunto de datos en el bucket de Cloud Storage

    (vm)$ gcloud storage cp ${SENTENCEPIECE_MODEL}.model ${STORAGE_BUCKET}
    (vm)$ gcloud storage cp ${SENTENCEPIECE_MODEL}.vocab ${STORAGE_BUCKET}
    (vm)$ gcloud storage cp tensorflow_datasets/wmt14_translate ${TFDS_DIR}/wmt14_translate --recursive
  4. Cambie al directorio de entrenamiento:

    (vm)$ cd /usr/share/tpu/models/
  5. Ejecuta la secuencia de comandos de entrenamiento:

    (vm)$ python3 official/nlp/train.py \
    --tpu=${TPU_NAME} \
    --experiment=wmt_transformer/large \
    --mode=train_and_eval \
    --model_dir=${MODEL_DIR} \
    --params_override="runtime.distribution_strategy=tpu, task.train_data.tfds_data_dir=${TFDS_DIR}, task.validation_data.tfds_data_dir=${TFDS_DIR}, task.sentencepiece_model_path=${SENTENCEPIECE_MODEL_PATH}, trainer.train_steps=10000, trainer.validation_interval=10000"

    Descripciones de las marcas de comandos

    tpu
    El nombre de Cloud TPU. Para ello, especifica la variable de entorno (TPU_NAME).
    experiment
    El modelo que se entrenará.
    mode
    Es el modo en el que se ejecutará la secuencia de comandos.
    model_dir
    El bucket de Cloud Storage en el que se almacenan los puntos de control y los resúmenes durante el entrenamiento. Puedes usar una carpeta existente para cargar puntos de control generados con anterioridad creados en una TPU del mismo tamaño y con la misma versión de TensorFlow.
    params_override
    Configura los parámetros del modelo.

De forma predeterminada, el modelo se evaluará después de cada 10,000 pasos. Para entrenar a la convergencia, cambia train_steps a 200,000. Puedes aumentar el la cantidad de pasos de entrenamiento o especificar con qué frecuencia ejecutar evaluaciones estos parámetros:

  • trainer.train_steps: Configura la cantidad total de pasos de entrenamiento que se deben ejecutar.
  • trainer.validation_interval: Establece la cantidad de pasos de entrenamiento que se deben ejecutar entre las evaluaciones.

El entrenamiento y la evaluación demoran aproximadamente 14 minutos en una v4-32 Cloud TPU Cuando se completa el entrenamiento y la evaluación, los mensajes similares a los siguientes:

I0209 22:19:49.143219 139751309618240 controller.py:310]  eval | step:  10000 | eval time:   73.6 sec | output:
    {'bleu_score': 19.401752948760986,
     'sacrebleu_score': 18.442741330886378,
     'validation_loss': 2.0558002}
 eval | step:  10000 | eval time:   73.6 sec | output:
    {'bleu_score': 19.401752948760986,
     'sacrebleu_score': 18.442741330886378,
     'validation_loss': 2.0558002}

Esta secuencia de comandos de entrenamiento entrena para 20,000 pasos y ejecuta la evaluación cada 2,000 pasos. Este entrenamiento y evaluación en particular toman aproximadamente 8 minutos en un pod de Cloud TPU v3-32. Cuando se completan el entrenamiento y la evaluación, aparece un mensaje similar al siguiente:

INFO:tensorflow:Writing to file /tmp/tmpdmlanxcf
I0218 21:09:19.100718 140509661046592 translate.py:184] Writing to file /tmp/tmpdmlanxcf
I0218 21:09:28.043537 140509661046592 transformer_main.py:118] Bleu score (uncased): 1.799112930893898
I0218 21:09:28.043911 140509661046592 transformer_main.py:119] Bleu score (cased): 1.730366237461567

Para entrenar hasta la convergencia, cambia train_steps a 200000. Tú puede aumentar la cantidad de pasos de entrenamiento o especificar la frecuencia de ejecución las evaluaciones estableciendo estos parámetros:

  • --train_steps: Configura la cantidad total de pasos de entrenamiento que se deben ejecutar.
  • --steps_between_evals: Cantidad de pasos de entrenamiento para ejecutar entre evaluaciones.

Cuando se completan el entrenamiento y la evaluación, aparece un mensaje similar al siguiente:

0509 00:27:59.984464 140553148962624 translate.py:184] Writing to file /tmp/tmp_rk3m8jp
I0509 00:28:11.189308 140553148962624 transformer_main.py:119] Bleu score (uncased): 1.3239131309092045
I0509 00:28:11.189623 140553148962624 transformer_main.py:120] Bleu score (cased): 1.2855342589318752

Limpia

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

  1. Desconéctate de la instancia de Compute Engine, si aún no lo hiciste:

    (vm)$ exit

    El mensaje ahora debería mostrar username@projectname, que indica que estás en Cloud Shell.

  2. Borra tus recursos de Cloud TPU y Compute Engine.

    $ gcloud compute tpus tpu-vm delete transformer-tutorial \
      --zone=us-central2-b
  3. Ejecuta gcloud CLI como se muestra y reemplaza bucket-name por el nombre del bucket de Cloud Storage que creaste para este instructivo:

    $ gcloud storage rm gs://bucket-name --recursive

¿Qué sigue?

Los instructivos de TensorFlow Cloud TPU suelen entrenar el modelo con un conjunto de datos de muestra. Los resultados de esta capacitación no se pueden usar para la inferencia. Para usar un modelo para inferencia, puedes entrenar los datos en un modelo conjunto de datos o uno propio. Modelos de TensorFlow entrenados con Cloud TPU generalmente requieren que los conjuntos de datos estén TFRecord.

Puedes usar la muestra de la herramienta de conversión de conjuntos de datos para convertir un conjunto de datos de clasificación de imágenes en formato TFRecord. Si no usas un modelo de clasificación de imágenes, deberás convertir tú mismo el conjunto de datos en formato TFRecord. Para obtener más información, consulta TFRecord y tf.Example

Ajuste de hiperparámetros

Para mejorar el rendimiento del modelo con tu conjunto de datos, puedes ajustar la hiperparámetros. Puedes encontrar información sobre los hiperparámetros comunes a todos Los modelos compatibles con TPU en GitHub: La información sobre los hiperparámetros específicos del modelo se puede encontrar en el código fuente de cada modelo. Para obtener más información sobre el ajuste de hiperparámetros, consulta Descripción general de ajuste de hiperparámetros y ajustes hiperparámetros.

Inferencia

Una vez que hayas entrenado tu modelo, puedes usarlo para inferencias (también denominado predicciones). Puedes usar la herramienta del convertidor de inferencia de Cloud TPU para preparar y optimizar un modelo de TensorFlow para la inferencia en Cloud TPU v5e. Para ver más información sobre la inferencia en Cloud TPU v5e, consulta Inferencia de Cloud TPU v5e introducción.