Entrenamiento previo de Wav2Vec2 en Cloud TPU con PyTorch

En este instructivo, se muestra cómo entrenar previamente el modelo Wav2Vec2 de FairSeq en un dispositivo de Cloud TPU con PyTorch. Puedes aplicar el mismo patrón a otros modelos de clasificación de imágenes optimizados con TPU que usen PyTorch y el conjunto de datos ImageNet.

El modelo de este instructivo se basa en el documento wav2vec 2.0: A Framework for Self-Monitored Learning of Speech Representations.

Objetivos

  • Crea y configura el entorno PyTorch.
  • Descarga datos de LibriSpeech de código abierto.
  • Ejecutar el trabajo de entrenamiento

Costos

En este instructivo, se usan 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. En la página del selector de proyectos de Google Cloud Console, selecciona o crea un proyecto de Google Cloud. Nota: Si no planeas conservar los recursos que creaste durante este procedimiento, crea un proyecto en lugar de seleccionar uno existente. Cuando termines, puedes borrar el proyecto y quitar todos los recursos asociados con él.
  2. Ve a la página del selector de proyectos y asegúrate de que la facturación esté habilitada para tu proyecto de Cloud. Aprende a confirmar que la facturación está habilitada para tu proyecto.
  1. Accede a tu cuenta de Google Cloud. Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.
  2. En la página del selector de proyectos de Google Cloud Console, selecciona o crea un proyecto de Google Cloud.

    Ir al selector de proyectos

  3. Asegúrate de que la facturación esté habilitada para tu proyecto de Cloud. Obtén información sobre cómo verificar si la facturación está habilitada en un proyecto.

  4. En la página del selector de proyectos de Google Cloud Console, selecciona o crea un proyecto de Google Cloud.

    Ir al selector de proyectos

  5. Asegúrate de que la facturación esté habilitada para tu proyecto de Cloud. Obtén información sobre cómo verificar si la facturación está habilitada en un proyecto.

  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.

Configura una instancia de Compute Engine

  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 la CLI de Google Cloud para usar el proyecto en el que deseas crear Cloud TPU.

    gcloud config set project ${PROJECT_ID}
    

    La primera vez que ejecutes este comando en una VM de Cloud Shell nueva, se mostrará la 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 GCP con tus credenciales.

  4. Desde Cloud Shell, inicia el recurso de Compute Engine necesario para este instructivo.

    gcloud compute instances create wav2vec2-tutorial \
      --zone=us-central1-a \
      --machine-type=n1-standard-64 \
      --image-family=torch-xla \
      --image-project=ml-images  \
      --boot-disk-size=200GB \
      --scopes=https://www.googleapis.com/auth/cloud-platform
    
  5. Conéctate a la instancia nueva de Compute Engine.

    gcloud compute ssh wav2vec2-tutorial --zone=us-central1-a
    

Inicia un recurso de Cloud TPU

  1. En la máquina virtual de Compute Engine, configura la versión de PyTorch.

    (vm) $ export PYTORCH_VERSION=1.11
    
  2. Inicia un recurso de Cloud TPU con el siguiente comando:

    (vm) $ gcloud compute tpus create wav2vec2-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-1.11 \
    --accelerator-type=v3-8
    
  3. Identifica la dirección IP del recurso de Cloud TPU ejecutando el siguiente comando:

    (vm) $ gcloud compute tpus describe wav2vec2-tutorial --zone=us-central1-a
    

    Este comando mostrará información sobre la TPU. Busca la línea con la etiqueta ipAddress. Reemplaza ip-address por la dirección IP en el siguiente comando:

    (vm) $ export TPU_IP_ADDRESS=ip-address
    (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    

Crea y configura el entorno PyTorch

  1. Inicia un entorno conda.

    (vm) $ conda activate torch-xla-1.11
    
  2. Configura variables de entorno para el recurso de Cloud TPU.

Descarga y prepara datos

Visita el sitio web de OpenSLR para ver los conjuntos de datos alternativos que puedes usar en esta tarea. En este instructivo, usaremos dev-clean.tar.gz, ya que tiene el tiempo de procesamiento previo más corto.

  1. Wav2Vec2 requiere algunas dependencias, instálalas ahora.

    (vm) $ pip install omegaconf hydra-core soundfile
    (vm) $ sudo apt-get install libsndfile-dev
    
  2. Descarga el conjunto de datos.

    (vm) $ curl https://www.openslr.org/resources/12/dev-clean.tar.gz --output dev-clean.tar.gz
    
  3. Extrae los archivos comprimidos. Tus archivos se almacenan en la carpeta LibriSpeech.

    (vm) $ tar xf dev-clean.tar.gz
    
  4. Descarga e instala el modelo de feria más reciente.

    (vm) $ git clone --recursive https://github.com/pytorch/fairseq.git
    (vm) $ cd fairseq
    (vm) $ pip install --editable .
    (vm) $ cd -
  5. Preparar el conjunto de datos Esta secuencia de comandos crea una carpeta llamada manifest/ con el puntero a los datos sin procesar (en LibriSpeech/).

    (vm) $ python fairseq/examples/wav2vec/wav2vec_manifest.py LibriSpeech/ --dest manifest/

Ejecuta el trabajo de entrenamiento

  1. Ejecuta el modelo con datos de LibriSpeech. Esta secuencia de comandos tarda alrededor de 2 horas en ejecutarse.

    (vm) $ fairseq-hydra-train \
    task.data=${HOME}/manifest \
    optimization.max_update=400000 \
    dataset.batch_size=4 \
    common.log_format=simple \
    --config-dir fairseq/examples/wav2vec/config/pretraining   \
    --config-name wav2vec2_large_librivox_tpu.yaml
    

Realiza una limpieza

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 user@projectname, que indica que estás en Cloud Shell.

  2. En Cloud Shell, usa la CLI de Google Cloud para borrar la instancia de VM de Compute Engine y la TPU:

    $ gcloud compute tpus execution-groups delete wav2vec2-tutorial --zone=us-central1-a
    

¿Qué sigue?

Escalamiento a los pods de Cloud TPU

Consulta el instructivo Entrena modelos de PyTorch en pods de Cloud TPU para escalar la tarea de entrenamiento previo en este instructivo a los potentes pods de Cloud TPU.

Prueba los siguientes colaboradores de PyTorch: