Pré-treinamento do Wav2Vec2 no Cloud TPU com PyTorch

Neste tutorial, mostramos como treinar previamente o modelo Wav2Vec2 de FairSeq em um dispositivo do Cloud TPU com o 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 documento wav2vec 2.0: A Framework for Self-supervisioned Learning of Speech Representations (em inglês).

Objetivos

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

Custos

Neste tutorial, usamos os seguintes componentes faturáveis do Google Cloud:

  • Compute Engine
  • Cloud TPU

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços. 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 do 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 que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover 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 o faturamento está ativado para o 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 verificar se o faturamento está ativado em um projeto.

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

  5. Verifique se o faturamento está ativado para seu projeto na nuvem. Saiba como verificar se o faturamento está ativado em um projeto.

  6. 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 CLI do Google Cloud para usar o projeto em que você quer criar a Cloud TPU.

    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.11
    
  2. Inicie um recurso do Cloud TPU usando o seguinte comando:

    (vm) $ gcloud compute tpus create wav2vec2-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-1.11 \
    --accelerator-type=v3-8
    
  3. Identifique o endereço IP para o recurso do Cloud TPU executando o comando a seguir:

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

    Esse comando exibirá informações sobre a TPU. Procure a linha marcada como ipAddress. Substitua ip-address pelo endereço IP no seguinte comando:

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

Crie e configure o ambiente PyTorch

  1. Inicie um ambiente conda.

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

Fazer o download e preparar dados

Acesse o site do OpenSLR para ver os conjuntos de dados alternativos que podem ser usados nesta tarefa. Neste tutorial, usaremos dev-clean.tar.gz porque ele tem o menor tempo de pré-processamento.

  1. O Wav2Vec2 requer 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 estão armazenados na pasta LibriSpeech.

    (vm) $ tar xf dev-clean.tar.gz
    
  4. Faça o download e instale o modelo 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. O script leva cerca de duas horas para ser executado.

    (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
    

Limpar

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 precisa ser user@projectname, mostrando que você está no Cloud Shell.

  2. No Cloud Shell, use a CLI do Google Cloud para excluir a instância de VM e a TPU do Compute Engine:

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

A seguir

Como escalonar para pods do Cloud TPU

Consulte o tutorial Como treinar modelos PyTorch em pods do Cloud TPU para escalonar a tarefa de pré-treinamento deste tutorial para os poderosos pods do Cloud TPU.

Teste as colabs do PyTorch: