Entrenamiento previo de Wav2Vec2 en Cloud TPU con PyTorch

En este instructivo, se muestra cómo entrenar el modelo Wav2Vec2 de FairSeq en un dispositivo 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: Un marco de trabajo para el aprendizaje automático supervisado de las representaciones de voz .

Objetivos

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

Costos

En este instructivo, se usan componentes facturables de Google Cloud, que incluyen los siguientes:

  • Compute Engine
  • Cloud TPU

Usa la calculadora de precios para generar una estimación de los costos según el uso previsto. Los usuarios nuevos de Google Cloud pueden ser elegibles 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 creados durante este procedimiento, crea un proyecto en lugar de seleccionar un proyecto existente. Cuando termines, puedes borrar el proyecto y quitar todos los recursos asociados con él.
  2. Ve a la página de selección de proyecto y asegúrate de que la facturación esté habilitada para tu proyecto de Cloud. Obtén información sobre cómo 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 proyecto

  3. Asegúrate de que la facturación esté habilitada para tu proyecto de Cloud. Descubre cómo confirmar que tienes habilitada la facturación en un proyecto.

  4. 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.

    Abrir Cloud Shell

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

    export PROJECT_ID=project-id
    
  3. Configura la herramienta de línea de comandos gcloud 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.8.1
    
  2. Inicia un recurso de Cloud TPU con el siguiente comando:

    (vm) $ gcloud compute tpus create w2v2-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-1.8 \
    --accelerator-type=v3-8
    
  3. Identifica la dirección IP para el recurso Cloud TPU.

    (vm) $ gcloud compute tpus list --zone=us-central1-a
    

Crea y configura el entorno PyTorch

  1. Inicia un entorno conda.

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

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

Descarga y prepara los 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 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 y, luego, instala el último modelo de Fairseq.

    (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 en los datos de libriSpeech. La secuencia de comandos tarda alrededor de 2 horas en ejecutarse.

    (vm) $ OMP_NUM_THREADS=1 python fairseq/train.py \
     manifest/ \
     --num-batch-buckets 3 \
     --tpu \
     --max-sentences 4 \
     --max-sentences-valid 4 \
     --required-batch-size-multiple 4 \
     --distributed-world-size 8 \
     --distributed-port 12597 \
     --update-freq 1 \
     --enable-padding \
     --log-interval 5 \
     --num-workers 6 \
     --task audio_pretraining \
     --criterion wav2vec \
     --arch wav2vec2 \
     --log-keys  "['prob_perplexity','code_perplexity','temp']" \
     --quantize-targets \
     --extractor-mode default \
     --conv-feature-layers '[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512,2,2)] * 2' \
     --final-dim 256 \
     --latent-vars 320 \
     --latent-groups 2 \
     --latent-temp '(2,0.5,0.999995)' \
     --infonce \
     --optimizer adam \
     --adam-betas '(0.9,0.98)' \
     --adam-eps 1e-06 \
     --lr-scheduler polynomial_decay \
     --total-num-update 400000 \
     --lr 0.0005 \
     --warmup-updates 32000 \
     --mask-length 10 \
     --mask-prob 0.65 \
     --mask-selection static \
     --mask-other 0 \
     --mask-channel-prob 0.1 \
     --encoder-layerdrop 0 \
     --dropout-input 0.0 \
     --dropout-features 0.0 \
     --feature-grad-mult 0.1 \
     --loss-weights '[0.1, 10]' \
     --conv-pos 128 \
     --conv-pos-groups 16 \
     --num-negatives 100 \
     --cross-sample-negatives 0 \
     --max-sample-size 250000 \
     --min-sample-size 32000 \
     --dropout 0.0 \
     --attention-dropout 0.0 \
     --weight-decay 0.01 \
     --max-tokens 1400000 \
     --max-epoch 10 \
     --save-interval 2 \
     --skip-invalid-size-inputs-valid-test \
     --ddp-backend no_c10d \
     --log-format simple

Realice 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 herramienta de línea de comandos de gcloud para borrar la instancia de VM de Compute Engine y la TPU:

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

¿Qué sigue?

Escala a pods de Cloud TPU

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

Prueba los siguientes colaboradores de PyTorch: