Como pré-treinar o Wav2Vec2 no Cloud TPU com PyTorch

Neste tutorial, mostramos como pré-treinar o modelo Wav2Vec2 da FairSeq em um dispositivo Cloud TPU com PyTorch. É possível aplicar o mesmo padrão a outros modelos de classificação de imagem otimizados para TPU que usam o PyTorch e o conjunto de dados do ImageNet.

O modelo deste tutorial é baseado no artigo wav2vec 2.0: A Framework for Learning autovisioned Learning of Speech presentations (em inglês).

Objetivos

  • Crie e configure o ambiente PyTorch.
  • Faça o download dos dados de código aberto do LibriSpeech.
  • Executar o job de treinamento.

Custos

Neste tutorial, há componentes faturáveis do Google Cloud, entre eles:

  • Compute Engine
  • Cloud TPU

Use a calculadora de preços para gerar uma estimativa de custos baseada na projeção de uso. Novos usuários do Google Cloud podem estar qualificados para uma avaliação gratuita.

Antes de começar

Antes de começar o tutorial, verifique se o projeto do Google Cloud foi configurado corretamente.

  1. No Console do Google Cloud, na página de seletor de projetos, selecione ou crie um projeto do Google Cloud. Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto em vez de selecionar um projeto existente. Depois de concluir essas etapas, é possível excluir o projeto, removendo todos os recursos associados a ele.
  2. Acesse a página do seletor de projetos e verifique se o faturamento está ativado para o projeto do Cloud. Saiba como confirmar se a cobrança está ativada para o seu projeto.
  1. Faça login na sua conta do Google Cloud. Se você começou a usar o Google Cloud agora, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.
  2. No Console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.

    Acessar o seletor de projetos

  3. Verifique se o faturamento está ativado para seu projeto na nuvem. Saiba como confirmar se o faturamento está ativado para o projeto.

  4. Este tutorial usa componentes faturáveis do Google Cloud. Consulte a página de preços da Cloud TPU para fazer uma estimativa dos custos. Para evitar cobranças desnecessárias, não se esqueça de apagar os recursos criados ao terminar de usá-los.

Configurar uma instância do Compute Engine

  1. Abra uma janela do Cloud Shell.

    Abra o Cloud Shell

  2. Crie uma variável para o ID do seu projeto.

    export PROJECT_ID=project-id
    
  3. Configure a ferramenta de linha de comando da gcloud para usar o projeto em que a Cloud TPU será criada.

    gcloud config set project ${PROJECT_ID}
    

    Na primeira vez que você executar esse comando em uma nova VM do Cloud Shell, será exibida uma página Authorize Cloud Shell. Clique em Authorize na parte inferior da página para permitir que gcloud faça chamadas de API do GCP com suas credenciais.

  4. No Cloud Shell, inicie o recurso do Compute Engine necessário para este tutorial.

    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. Conecte-se à nova instância do Compute Engine.

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

Iniciar um recurso da Cloud TPU

  1. Na máquina virtual do Compute Engine, defina a versão do PyTorch.

    (vm) $ export PYTORCH_VERSION=1.9
    
  2. Inicie um recurso do Cloud TPU usando o seguinte comando:

    (vm) $ gcloud compute tpus create w2v2-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-1.9 \
    --accelerator-type=v3-8
    
  3. Identifique o endereço IP do recurso do Cloud TPU.

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

Crie e configure o ambiente PyTorch

  1. Inicie um ambiente conda.

    (vm) $ conda activate torch-xla-1.9
    
  2. Configure variáveis de ambiente para o recurso Cloud TPU.

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

Fazer o download dos dados e prepará-los

Visite o site do OpenSLR para ver os conjuntos de dados alternativos que você pode usar nessa tarefa. Neste tutorial, usaremos dev-clean.tar.gz porque ele tem o menor tempo de pré-processamento.

  1. O Wav2Vec2 exige algumas dependências. Instale-as agora.

    (vm) $ pip install omegaconf hydra-core soundfile
    (vm) $ sudo apt-get install libsndfile-dev
    
  2. Faça o download do conjunto de dados.

    (vm) $ curl https://www.openslr.org/resources/12/dev-clean.tar.gz --output dev-clean.tar.gz
    
  3. Extraia os arquivos compactados. Seus arquivos são armazenados na pasta LibriSpeech.

    (vm) $ tar xf dev-clean.tar.gz
    
  4. Faça o download e instale o modelo do Fairseq mais recente.

    (vm) $ git clone --recursive https://github.com/pytorch/fairseq.git
    (vm) $ cd fairseq
    (vm) $ pip install --editable .
    (vm) $ cd -
  5. Preparar o conjunto de dados. Esse script cria uma pasta chamada manifest/ com o ponteiro para os dados brutos (em LibriSpeech/).

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

Executar o job de treinamento

  1. Execute o modelo nos dados do LibriSpeech. A execução deste script leva cerca de duas horas.

    (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

Limpeza

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

  1. Desconecte-se da instância do Compute Engine, caso ainda não tenha feito isso:

    (vm)$ exit
    

    Agora, o prompt será user@projectname, mostrando que você está no Cloud Shell.

  2. No Cloud Shell, use a ferramenta de linha de comando gcloud para excluir a instância de VM do Compute Engine e a TPU:

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

A seguir

Como escalonar para os pods do Cloud TPU

Consulte o tutorial Como treinar modelos PyTorch em pods do Cloud TPU para escalonar a tarefa de pré-treinamento neste tutorial para os avançados pods do Cloud TPU.

Teste as colabs do PyTorch: