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 e usados para acelerar as cargas de trabalho de machine learning. É possível executar jobs de treinamento no AI Platform com o Cloud TPU. O AI Platform fornece uma interface de gerenciamento de tarefas 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 (em inglês), seus modelos poderão ser executados no hardware do 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 do Cloud TPU associada ao seu projeto do GCP:

  1. Chame projects.getConfig para receber o nome da conta de serviço do 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 do Cloud TPU como membro do 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 gcloud:

Console

  1. Faça login no Console do Google Cloud Platform (em inglês) 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 do Cloud TPU:

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

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

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

Executar o modelo de amostra ResNet-50

Nesta seção, mostramos como treinar a referência Modelo do TensorFlow ResNet-50, usando um conjunto de dados falso fornecido 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 comandos a seguir 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.14.tar.gz
    tar -xzvf r1.14.tar.gz && rm r1.14.tar.gz
    
  2. Acesse o diretório official na estrutura de diretório descompactada:

    cd tpu-r1.14/models/official/
    
  3. Edite ./resnet/resnet_main.py e altere o código para usar importações relativas explícitas ao importar submódulos. Por exemplo, altere o seguinte:

    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 regiões a seguir permitem 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.14 \
            --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 no 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 regiões a seguir permitem 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 previsão on-line/em lote, leia o guia para regiões.

Controle de versões do TensorFlow e AI Platform

As versões 1.12, 1.13 e 1.14 do ambiente de execução do AI Platform estão disponíveis para treinar os modelos no 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 do 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, use o TPUClusterResolver (em inglês) para se conectar ao servidor gRPC em execução na VM da TPU. O TPUClusterResolver retorna o endereço IP e a porta do Cloud TPU.

O exemplo a seguir mostra como o código de amostra do ResNet-50 (em inglês) 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 Estimator do TensorFlow (ambos em inglês) de alto nível.

  • 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 a partir do ponto em que parou.

Veja a lista de operações do TensorFlow disponíveis no 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 ativada para TPU de uma destas três maneiras:

  • Use o nível de escalonamento BASIC_TPU. Use esse método para acessar os aceleradores da TPU v2.
  • Use um worker cloud_tpu e um tipo de máquina legada para a VM principal. Use esse método para acessar os aceleradores da TPU v2.
  • Use um worker cloud_tpu e um tipo de máquina do Compute Engine para a VM principal. É possível usar esse método para acessar os aceleradores TPU v2 ou TPU v3. Os aceleradores da TPU v3 estão disponíveis na versão Beta.

Máquina ativada para TPU básica

Defina o nível de escalonamento como BASIC_TPU para conseguir uma VM principal e uma VM de TPU, incluindo uma TPU com oito núcleos TPU v2, como você fez ao executar a amostra anterior.

Worker da TPU em uma configuração de tipo de máquina legada

Como alternativa, defina 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 principal para usar um tipo de máquina legada adequado aos requisitos do job.
  • Configure workerType como cloud_tpu para conseguir uma VM de TPU, incluindo um Cloud TPU com oito núcleos TPU v2.
  • Defina workerCount como 1.
  • Não especifique um servidor de parâmetros ao usar um Cloud TPU. O serviço rejeitará a solicitação de job se parameterServerCount for maior que zero.

No exemplo a seguir, mostramos um arquivo config.yaml que usa esse tipo de configuração:

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

Worker da 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.

É possível usar esse tipo de configuração para configurar um worker da TPU com oito núcleos TPU v2 (semelhante a uma configuração sem acceleratorConfig) ou um worker da TPU com oito núcleos TPU v3 (Beta). Leia mais sobre a diferença entre os aceleradores TPU v2 e TPU v3.

Usar um tipo de máquina do Compute Engine também oferece mais flexibilidade para configurar sua VM mestre:

  • 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 requisitos do job.
  • Defina workerType como cloud_tpu.
  • Adicione um workerConfig com um campo acceleratorConfig. Dentro de acceleratorConfig, defina type como TPU_V2 ou TPU_V3 e count como 8. Não é possível anexar qualquer outro número de núcleos de TPU.
  • Defina workerCount como 1.
  • Não especifique um servidor de parâmetros ao usar um Cloud TPU. O serviço rejeitará a solicitação de job se parameterServerCount for maior que zero.

No exemplo a seguir, mostramos um arquivo config.yaml que usa esse tipo de configuração:

TPU v2

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

TPU v3 (Beta)

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

Como usar um contêiner personalizado em um worker da TPU

Se você quiser executar um contêiner personalizado em seu worker da TPU em vez de usar uma das versões do ambiente de execução do AI Platform compatíveis com TPUs, especifique mais um campo de configuração ao enviar seu job de treinamento. Defina tpuTfVersion como uma versão do ambiente de execução que inclua a versão do TensorFlow que o contêiner usa. Especifique uma versão do ambiente de execução compatível para treinamento com TPUs.

Como você está configurando o job para usar um contêiner personalizado, o AI Platform Training não usa o ambiente de execução ao executar o job de treinamento. No entanto, ele precisa que esse campo seja preenchido para preparar adequadamente o worker da TPU para a versão do TensorFlow usada pelo contêiner personalizado.

No exemplo a seguir, mostramos um arquivo config.yaml com uma configuração de TPU semelhante àquela da seção anterior, exceto que neste caso a VM principal e o worker da TPU executam contêineres personalizados diferentes:

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/your-project-id/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/your-project-id/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 1.14
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3 (Beta)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/your-project-id/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/your-project-id/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 1.14
    acceleratorConfig:
      type: TPU_V3
      count: 8

Se você usar o comando gcloud beta ai-platform jobs submit training para enviar seu job de treinamento, poderá especificar o campo da API tpuTfVersion com a sinalização --tpu-tf-version, em vez de em um arquivo config.yaml.

Saiba mais sobre o treinamento distribuído com contêineres personalizados.

A seguir

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

Enviar comentários sobre…

Precisa de ajuda? Acesse nossa página de suporte.