Como usar TPUs para treinar seu modelo

As unidades de processamento de tensor (TPUs, na sigla em inglês) são ASICs desenvolvidos pelo Google usados para acelerar as cargas de trabalho de machine learning. É possível executar seus jobs de treinamento no AI Platform por meio da Cloud TPU. O AI Platform oferece uma interface de gerenciamento de jobs para que você não precise gerenciar a TPU por conta própria. Use a API jobs do AI Platform da mesma forma que a utiliza no treinamento em uma CPU ou GPU.

Com as APIs de alto nível do TensorFlow, seus modelos poderão ser executados no hardware da Cloud TPU.

Configurar e testar o ambiente do GCP

Configure seu ambiente do GCP seguindo as instruções da seção de configuração do guia de primeiros passos.

Autorize o acesso da Cloud TPU ao projeto

Siga estas etapas para autorizar o nome da conta de serviço da Cloud TPU associada ao seu projeto do GCP:

  1. Chame projects.getConfig para receber o nome da conta de serviço da Cloud TPU. Exemplo:

    curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
        https://ml.googleapis.com/v1/projects/<your-project-id>:getConfig
    
  2. Salve o valor do campo tpuServiceAccount retornado pela API.

Agora adicione a conta de serviço da Cloud TPU como membro do seu projeto, com o papel Agente de serviço do Cloud ML. Conclua as etapas a seguir no Console do Google Cloud Platform ou use o comando da gcloud:

Console

  1. Faça login no Console do Google Cloud Platform e escolha o projeto em que a TPU está sendo usada.
  2. Escolha IAM e administrador > IAM.
  3. Clique no botão Adicionar para adicionar um membro ao projeto.
  4. Digite a conta de serviço de TPU na caixa de texto Membros.
  5. Clique na lista suspensa Papéis.
  6. Ative o papel Agente de serviço do Cloud ML (Gerenciamento de serviços > Agente de serviço do Cloud ML).

gcloud

  1. Defina as variáveis de ambiente que contêm o ID do projeto e a conta de serviço da Cloud TPU:

    PROJECT_ID=your-project-id
    SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
    
  2. Conceda o papel ml.serviceAgent à conta de serviço da Cloud TPU:

    gcloud projects add-iam-policy-binding $PROJECT_ID \
        --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
    

Para ver mais detalhes sobre a concessão de papéis a contas de serviço, consulte a documentação do Cloud IAM.

Execute o modelo de amostra ResNet-50

Nesta seção, você verá como treinar o modelo de referência Tensorflow ResNet-50 usando um conjunto de dados fictício disponível em gs://cloud-tpu-test-datasets/fake_imagenet. No exemplo de job, o nível de escalonamento predefinido BASIC_TPU é usado para a configuração da sua máquina. As seções posteriores do guia contêm informações sobre como definir uma configuração personalizada.

Execute os seguintes comandos para receber o código e enviar seu job de treinamento no AI Platform:

  1. Faça o download do código para o modelo de referência:

    mkdir tpu-demos && cd tpu-demos
    wget https://github.com/tensorflow/tpu/archive/r1.12.tar.gz
    tar -xzvf r1.12.tar.gz && rm r1.12.tar.gz
    
  2. Acesse o diretório official da estrutura de diretórios descompactada:

    cd tpu-r1.12/models/official/
    
  3. Edite ./resnet/resnet_main.py e altere o código para usar importações relativas explícitas na importação de submódulos. Por exemplo, mude:

    import resnet_model
    

    Para:

    from . import resnet_model
    

    Use o padrão acima para alterar todas as outras importações no arquivo e salve o arquivo.

  4. Verifique se há importações de submódulos em outros arquivos da amostra e atualize-as para usar importações relativas explícitas também.

  5. Configure algumas variáveis de ambiente:

    JOB_NAME=tpu_1
    STAGING_BUCKET=gs://my_bucket_for_staging
    REGION=us-central1
    DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
    OUTPUT_PATH=gs://my_bucket_for_model_output
    

    Atualmente, as seguintes regiões dão acesso a TPUs:

    • us-central1

  6. Envie seu job de treinamento usando o comando gcloud ai-platform jobs submit training :

    gcloud ai-platform jobs submit training $JOB_NAME \
            --staging-bucket $STAGING_BUCKET \
            --runtime-version 1.12 \
            --scale-tier BASIC_TPU \
            --module-name resnet.resnet_main \
            --package-path resnet/ \
            --region $REGION \
            -- \
            --data_dir=$DATA_DIR \
            --model_dir=$OUTPUT_PATH
    

Mais detalhes sobre como treinar um modelo na Cloud TPU

Na parte inicial deste guia, você aprendeu a usar o código de amostra do ResNet-50. Nesta seção, você verá mais informações sobre como configurar um job e treinar um modelo no AI Platform com a Cloud TPU.

Como especificar uma região que oferece TPUs

Você precisa executar seu job em uma região que tenha TPUs disponíveis. Atualmente, as seguintes regiões dão acesso a TPUs:

  • us-central1

Para entender completamente as regiões disponíveis para serviços do AI Platform, incluindo treinamento de modelos e predição on-line/em lote, leia o guia para regiões.

Controle de versões do TensorFlow e AI Platform

As versões 1.11 e 1.12 do ambiente de execução do AI Platform estão disponíveis para treinar os modelos na Cloud TPU. Veja mais sobre as versões do ambiente de execução do AI Platform e sobre as versões correspondentes do TensorFlow.

A política de controle de versões é a mesma da Cloud TPU. Na sua solicitação de job de treinamento, especifique uma versão de ambiente de execução que esteja disponível para TPUs e corresponda à versão do TensorFlow usada no seu código de treinamento.

Como se conectar com o servidor gRPC da TPU

No programa do Tensorflow, você deverá usar o TPUClusterResolver para se conectar ao servidor gRPC em execução na VM da TPU. O TPUClusterResolver retorna o endereço IP e a porta da Cloud TPU.

O exemplo a seguir mostra como o código de amostra do ResNet-50 usa o TPUClusterResolver:

tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
    FLAGS.tpu,
    zone=FLAGS.tpu_zone,
    project=FLAGS.gcp_project)

config = tpu_config.RunConfig(
    cluster=tpu_cluster_resolver,
    model_dir=FLAGS.model_dir,
    save_checkpoints_steps=max(600, FLAGS.iterations_per_loop),
    tpu_config=tpu_config.TPUConfig(
        iterations_per_loop=FLAGS.iterations_per_loop,
        num_shards=FLAGS.num_cores,
        per_host_input_for_training=tpu_config.InputPipelineConfig.PER_HOST_V2))  # pylint: disable=line-too-long

Como atribuir operações a TPUs

Para usar as TPUs em uma máquina, use a API TPUEstimator do TensorFlow, herdada da API de alto nível Estimator.

  • A TPUEstimator processa vários detalhes da execução em dispositivos TPU, como réplica de entradas e modelos para cada núcleo e retorno ao host periodicamente para executar hooks.
  • A API de alto nível TensorFlow oferece também muitas outras facilidades. Em particular, a API salva e restaura os checkpoints do modelo para que você possa retomar um job de treinamento interrompido no ponto em que parou.

Veja a lista de operações do TensorFlow disponíveis na Cloud TPU.

Como configurar uma máquina personalizada de TPU

Um job de treinamento de TPU é executado em uma configuração de duas VMs. Uma VM (mestre) executa seu código Python. A mestre controla o servidor do TensorFlow em execução em um worker da TPU.

Para usar uma TPU com o AI Platform, configure seu job de treinamento para acessar uma máquina habilitada para TPU de uma destas três maneiras:

  • Use o nível de escalonamento BASIC_TPU.
  • Use um worker cloud_tpu e um tipo de máquina do AI Platform para a VM mestre.
  • Use um worker cloud_tpu e um tipo de máquina do Compute Engine para a VM principal.

Máquina habilitada para TPU básica

Defina o nível de escalonamento como BASIC_TPU para ter uma VM mestre e uma VM de TPU, incluindo uma TPU, igual à amostra executada acima.

Worker TPU em uma configuração de tipo de máquina do AI Platform

Como alternativa, você pode definir uma configuração de máquina personalizada, caso precise de mais recursos de computação na VM mestre:

  • Defina o nível de escalonamento como CUSTOM.
  • Configure a VM mestre para usar um tipo de máquina do AI Platform que atenda aos seus requisitos de trabalho.
  • Defina workerType como cloud_tpu para obter uma VM de TPU, incluindo uma Cloud TPU.
  • Defina workerCount como 1.
  • Não especifique um servidor de parâmetros ao usar uma Cloud TPU. O serviço rejeitará a solicitação de job se parameterServerCount for maior que zero.

O exemplo a seguir mostra um config.yaml que usa esse tipo de configuração:

trainingInput:
  scaleTier: CUSTOM
  masterType: complex_model_m
  workerType: cloud_tpu
  workerCount: 1

Worker TPU em uma configuração de tipo de máquina do Compute Engine

Também é possível definir uma configuração de máquina personalizada com um tipo de máquina do Compute Engine para sua VM mestre e um acceleratorConfig anexado à sua VM de TPU.

Atualmente, isso fornece os mesmos recursos de TPU (oito núcleos de TPU v2) que uma configuração sem um acceleratorConfig. No entanto, usar um tipo de máquina do Compute Engine pode oferecer mais flexibilidade para configurar sua VM mestre. Esse tipo de configuração de job de treinamento está no estágio de lançamento Beta:

  • Defina o nível de escalonamento como CUSTOM.
  • Configure a VM mestre para usar um tipo de máquina do Compute Engine que atenda aos seus requisitos de trabalho.
  • Defina workerType como cloud_tpu.
  • Adicione um workerConfig com um campo acceleratorConfig. Dentro desse acceleratorConfig, defina type como TPU_V2, e count como 8. Não é possível anexar qualquer outro número de núcleos de TPU v2.
  • Defina workerCount como 1.
  • Não especifique um servidor de parâmetros ao usar uma Cloud TPU. O serviço rejeitará a solicitação de job se parameterServerCount for maior que zero.

O exemplo a seguir mostra um config.yaml que usa esse tipo de configuração:

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerConfig:
    acceleratorConfig:
      type: TPU_V2
      count: 8
  workerCount: 1

A seguir

Esta página foi útil? Conte sua opinião sobre:

Enviar comentários sobre…

AI Platform para TensorFlow