Treinamento do RetinaNet na Cloud TPU (TF 2.x)

Neste documento, descrevemos uma implementação do modelo de detecção de objetos do RetinaNet. O código está disponível no GitHub.

Nas instruções abaixo, presume-se que você conheça a execução de um modelo na Cloud TPU. Se você conhece a Cloud TPU há pouco tempo, consulte o guia de início rápido para ver os conceitos básicos.

Se você planeja treinar em uma fração do pod da TPU, veja como realizar o treinamento em pods de TPU para entender as alterações de parâmetros necessárias para frações de pod.

Objetivos

  • Criar um bucket do Cloud Storage para armazenar o conjunto de dados e a saída do modelo.
  • Preparar o conjunto de dados COCO.
  • Configurar uma VM do Compute Engine e um nó da Cloud TPU para treinamento e avaliação.
  • Executar treinamento e avaliação em uma única Cloud TPU ou em um pod da Cloud TPU.

Custos

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

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

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

Nesta seção, você verá como configurar o bucket do Cloud Storage e a VM 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 gcloud para usar o projeto em que a Cloud TPU será criada.

    gcloud config set project ${PROJECT_ID}
    
  4. Crie uma conta de serviço para o projeto da Cloud TPU.

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
    

    O comando retorna uma conta de serviço do Cloud TPU com o formato a seguir:

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

  5. Crie um bucket do Cloud Storage usando o seguinte comando:

    gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 -b on gs://bucket-name
    

    Esse bucket do Cloud Storage armazena os dados usados para treinar o modelo e os resultados do treinamento. A ferramenta ctpu up usada neste tutorial configura permissões padrão para a conta de serviço do Cloud TPU configurada na etapa anterior. Caso queira permissões mais específicas, consulte as permissões de nível de acesso.

  6. Inicie uma VM do Compute Engine usando o comando ctpu up.

    $ ctpu up --vm-only \
    --name=retinanet-tutorial \
    --disk-size-gb=300 \
    --machine-type=n1-standard-8 \
    --zone=europe-west4-a \
    --tf-version=2.3.1
    

    Descrições de sinalizações de comando

    vm-only
    Cria a VM sem criar uma Cloud TPU. Por padrão, o comando ctpu up cria uma VM e uma Cloud TPU.
    name
    O nome da Cloud TPU a ser criada.
    disk-size-gb
    O tamanho do disco rígido em GB da VM criada pelo comando ctpu up.
    machine-type
    O tipo de máquina da VM do Compute Engine a ser criado.
    zone
    A zona em que você planeja criar a Cloud TPU.
    tf-version
    A versão do Tensorflow ctpu é instalada na VM.

    Para mais informações sobre o utilitário do CTPU, consulte a Referência do CTPU.

  7. Quando solicitado, pressione y para criar os recursos da Cloud TPU.

    Quando o comando ctpu up terminar a execução, verifique se o prompt do shell foi alterado de username@projectname para username@vm-name. Essa alteração mostra que você fez login na VM do Compute Engine.

    gcloud compute ssh retinanet-tutorial --zone=europe-west4-a
    

    Ao continuar essas instruções, execute cada comando que começa com (vm)$ na instância do Compute Engine.

  8. Instalar pacotes extras

    O aplicativo de treinamento do RetinaNet requer vários pacotes extras. Instale-os agora:

    (vm)$ sudo apt-get install -y python3-tk
    (vm)$ pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow
    
    (vm)$ pip3 install --user 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI'
    
    (vm)$ sudo pip3 install --user -r /usr/share/models/official/requirements.txt
    

Preparar o conjunto de dados COCO

O conjunto de dados COCO será armazenado no Cloud Storage. Portanto, defina uma variável de bucket de armazenamento especificando o nome do bucket que você criou:

(vm)$ export STORAGE_BUCKET=gs://bucket-name
(vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco

Execute o script download_and_preprocess_coco.sh para converter o conjunto de dados COCO em um conjunto de TFRecords (*.tfrecord) esperado pelo aplicativo de treinamento.

(vm)$ sudo bash /usr/share/tpu/tools/datasets/download_and_preprocess_coco.sh ./data/dir/coco

Isso instala as bibliotecas necessárias e executa o script de pré-processamento. A saída são alguns arquivos *.tfrecord no diretório de dados. O script de download e conversão do COCO leva aproximadamente 1 hora para ser concluído.

Copiar os dados para o bucket do Cloud Storage

Depois de converter os dados para TFRecords, use o comando gsutil para copiá-los do armazenamento local para o bucket do Cloud Storage. Também é preciso copiar os arquivos de anotação. Eles ajudam a validar o desempenho do modelo.

(vm)$ gsutil -m cp ./data/dir/coco/*.tfrecord ${DATA_DIR}
(vm)$ gsutil cp ./data/dir/coco/raw-data/annotations/*.json ${DATA_DIR}

Configurar o ambiente de treinamento

  1. Crie uma Cloud TPU usando o comando ctpu up.

    (vm)$ ctpu up --tpu-only \
      --tpu-size=v3-8 \
      --zone=europe-west4-a \
      --name=retinanet-tutorial \
      --tf-version=2.3.1
    

    Para mais informações sobre o utilitário do CTPU, consulte a Referência do CTPU.

    Descrições de sinalizações de comando

    tpu-only
    Crie apenas uma Cloud TPU. Por padrão, o comando ctpu up cria uma VM e uma Cloud TPU.
    tpu-size
    O tipo da Cloud TPU a ser criada.
    zone
    A zona em que você planeja criar a Cloud TPU.
    name
    O nome da Cloud TPU a ser criada.
    tf-version
    A versão do ctpu do Tensorflow é instalada na VM.
    Parâmetro Descrição

    Para mais informações sobre o utilitário do CTPU, consulte a Referência do CTPU.

  2. Quando solicitado, pressione y para criar os recursos da Cloud TPU.

    Quando o comando ctpu up terminar a execução, verifique se o prompt do shell foi alterado de username@projectname para username@vm-name. Essa alteração mostra que você fez login na VM do Compute Engine.

    gcloud compute ssh retinanet-tutorial --zone=europe-west4-a
    

    Ao continuar essas instruções, execute cada comando que começa com (vm)$ na instância do Compute Engine.

  3. Crie uma variável de ambiente para o nome da TPU.

    (vm)$ export TPU_NAME=retinanet-tutorial
    
  4. Adicione a pasta de nível superior /models ao caminho do Python com o comando

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
    

Treinamento sobre dispositivos de Cloud TPU único

Os scripts de treinamento a seguir foram executados em uma Cloud TPU v3-8. Isso levará mais tempo, mas também é possível executá-los em um Cloud TPU v2-8.

Este exemplo de script abaixo é treinado por apenas 10 etapas e leva menos de cinco minutos para ser executado em um nó de TPU v3-8. O treinamento para a conversão leva cerca de 22.500 etapas e aproximadamente 1 hora e meia em uma TPU v3-8 do Cloud TPU.

  1. Configure as seguintes variáveis de ambiente:

    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/retinanet-train
    (vm)$ export RESNET_CHECKPOINT=gs://cloud-tpu-checkpoints/retinanet/resnet50-checkpoint-2018-02-07
    (vm)$ export TRAIN_FILE_PATTERN=${DATA_DIR}/train-*
    (vm)$ export EVAL_FILE_PATTERN=${DATA_DIR}/val-*
    (vm)$ export VAL_JSON_FILE=${DATA_DIR}/instances_val2017.json
    
  2. Execute o script de treinamento:

    (vm)$ python3 /usr/share/models/official/vision/detection/main.py \
         --strategy_type=tpu \
         --tpu=${TPU_NAME} \
         --model_dir=${MODEL_DIR} \
         --mode="train" \
         --params_override="{ type: retinanet, train: { total_steps: 10, checkpoint: { path: ${RESNET_CHECKPOINT}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN} }, eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: 5000 } }"
    

    Descrições de sinalizações de comando

    strategy_type
    Para treinar o modelo RetinaNet em uma TPU, é necessário definir o distribution_strategy como tpu.
    tpu
    O nome da Cloud TPU. Isso é definido usando a variável de ambiente TPU_NAME.
    model_dir
    O diretório em que os pontos de controle e os resumos são armazenados durante o treinamento do modelo. Se a pasta não existir, o programa criará uma. Ao usar uma Cloud TPU, o model_dir precisa ser um caminho do Cloud Storage (gs://...). É possível reutilizar uma pasta existente para carregar dados do ponto de controle atuais e armazenar outros pontos, contanto que os anteriores tenham sido criados usando uma Cloud TPU com o mesmo tamanho e versão do TensorFlow.
    mode
    Defina isso como train para treinar o modelo ou eval para avaliar o modelo.
    params_override
    Uma string JSON que modifica parâmetros de script padrão. Para mais informações sobre parâmetros de script, consulte /usr/share/models/official/vision/detection/main.py.

Enquanto o modelo está em treinamento, é possível ver o progresso visualizando a saída do registro. Um resultado semelhante a este mostra que o treinamento está progredindo normalmente:

31517803669, 'learning_rate': 0.08, 'box_loss': 0.0006472870009019971,
'l2_regularization_loss': 0.09328179806470871}
I0210 21:59:19.888985 139927795508992 distributed_executor.py:49]
Saving model as TF checkpoint: gs://bucket-eu/retinanet-model/ctl_step_2500.ckpt-5
I0210 22:01:07.714749 139927795508992 distributed_executor.py:446]
Train Step: 3000/22500  / loss = {'model_loss': 0.08362223953008652,
'total_loss': 0.17120523750782013, 'cls_loss': 0.057121846824884415,
'learning_rate': 0.08, 'box_loss': 0.0005300078773871064,
'l2_regularization_loss': 0.08758299797773361} / training metric =
{'model_loss': 0.08362223953008652, 'total_loss': 0.17120523750782013, 'cls_loss': 0.0
57121846824884415, 'learning_rate': 0.08, 'box_loss': 0.0005300078773871064,
'l2_regularization_loss': 0.08758299797773361}
I0210 22:01:15.813422 139927795508992 distributed_executor.py:49]
Saving model as TF checkpoint: gs://bucket-eu/retinanet-model/ctl_step_3000.ckpt-6

Avaliação de dispositivos de Cloud TPU único

O procedimento a seguir usa os dados de avaliação do COCO. Ele leva cerca de 10 minutos para ser executado nas etapas de avaliação.

  1. Configure as seguintes variáveis de ambiente:

    (vm)$ export EVAL_SAMPLES=5000
    
  2. Execute o script de avaliação:

    (vm)$ python3 /usr/share/models/official/vision/detection/main.py \
          --strategy_type=tpu \
          --tpu=${TPU_NAME} \
          --model_dir=${MODEL_DIR} \
          --mode="eval" \
          --params_override="{ type: retinanet, eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: ${EVAL_SAMPLES} } }"
    

    Descrições de sinalizações de comando

    strategy_type
    Para treinar o modelo RetinaNet em uma TPU, é necessário definir o distribution_strategy como tpu.
    tpu
    O nome da Cloud TPU. Isso é definido usando a variável de ambiente TPU_NAME.
    model_dir
    O diretório em que os pontos de controle e os resumos são armazenados durante o treinamento do modelo. Se a pasta não existir, o programa criará uma. Ao usar uma Cloud TPU, o model_dir precisa ser um caminho do Cloud Storage (gs://...). É possível reutilizar uma pasta existente para carregar dados do ponto de controle atuais e armazenar outros pontos, contanto que os anteriores tenham sido criados usando uma Cloud TPU com o mesmo tamanho e versão do TensorFlow.
    mode
    Defina model como eval para avaliar o modelo.
    params_override
    Uma string JSON que modifica parâmetros de script padrão. Para mais informações sobre parâmetros de script, consulte /usr/share/models/official/vision/detection/main.py.

    No final da avaliação, você verá mensagens semelhantes às seguintes no console:

    Accumulating evaluation results...
    DONE (t=7.66s).
     Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.000
     Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.000
     Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.000
     Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
     Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
     Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
     Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000
    

    A partir daqui, é possível concluir este tutorial e limpar os recursos do GCP ou explorar a execução do modelo em pods da Cloud TPU.

Como escalonar o modelo com os pods da Cloud TPU

Para resultados mais rápidos, escalone o modelo com os pods da Cloud TPU. O modelo RetinaNet totalmente compatível pode funcionar com a fração de pod v2-32:

  1. Exclua o recurso da Cloud TPU criado para treinar o modelo em um único dispositivo.

    (vm)$ ctpu delete --tpu-only --zone=europe-west4-a --name=retinanet-tutorial
  2. Execute o comando ctpu up, usando o parâmetro tpu-size para especificar a fração de pod a ser usada. Por exemplo, o comando a seguir usa uma fração de pod v2-32.

    (vm)$ ctpu up --tpu-only \
    --tpu-size=v2-32  \
    --zone=europe-west4-a \
    --name=retinanet-tutorial \
    --tf-version=2.3.1

    Descrições de sinalizações de comando

    tpu-only
    Cria apenas a Cloud TPU. Por padrão, o comando ctpu up cria uma VM e uma Cloud TPU.
    tpu-size
    O tipo da Cloud TPU a ser criada.
    zone
    A zona em que você planeja criar a Cloud TPU.
    name
    O nome da Cloud TPU a ser criada.
    tf-version
    A versão do ctpu do Tensorflow é instalada na VM.

    Para mais informações sobre o utilitário do CTPU, consulte a Referência do CTPU.

  3. Defina a variável de nome da Cloud TPU. Este será um nome definido com o parâmetro --name ou o padrão, seu nome de usuário:

    (vm)$ export TPU_NAME=retinanet-tutorial
    
  4. Configure a seguinte variável de ambiente:

    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/retinanet-pod
    
  5. Executar o script de treinamento do pod em um nó da TPU v2-32

    O script de treinamento de amostra a seguir foi executado em um pod v2-32 do Cloud TPU. Ele é treinado por apenas 10 etapas e leva menos de cinco minutos para ser executado. Para treinar para a conversão, são necessários 219 passos e leva aproximadamente 50 minutos em um pod de TPU v2-32.

    (vm)$  python3 /usr/share/models/official/vision/detection/main.py \
        --strategy_type=tpu \
        --tpu=${TPU_NAME} \
        --model_dir=${MODEL_DIR} \
        --mode="train" \
        --params_override="{ type: retinanet, train: { total_steps: 10, batch_size: 256, checkpoint: { path: ${RESNET_CHECKPOINT}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN} }, eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: 5000 } }"
    

    Descrições de sinalizações de comando

    tpu
    Especifica o nome da Cloud TPU. Isso é definido usando a variável de ambiente TPU_NAME.
    model_dir
    Especifica o diretório em que os pontos de controle e os resumos são armazenados durante o treinamento do modelo. Se a pasta não existir, o programa criará uma. Ao usar uma Cloud TPU, o model_dir precisa ser um caminho do Cloud Storage (gs://...). É possível reutilizar uma pasta existente para carregar dados dos pontos de controle atuais e armazenar outros pontos, contanto que os anteriores tenham sido criados usando uma Cloud TPU com o mesmo tamanho e versão do TensorFlow.

Limpar

Para evitar cobranças dos recursos usados neste tutorial na conta do Google Cloud Platform:

  1. Encerre a conexão com a VM do Compute Engine:

    (vm)$ exit
    

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

  2. Na VM ou no Cloud Shell, use o seguinte comando para excluir a VM e o Cloud TPU:

    $ ctpu delete --name=retinanet-tutorial \
      --zone=europe-west4-a
    
  3. Execute ctpu status especificando sua zona para garantir que você não tenha instâncias alocadas para evitar cobranças desnecessárias pelo uso da TPU. A exclusão pode levar vários minutos. Uma resposta como esta indica que não há mais instâncias alocadas:

    $ ctpu status --name=retinanet-tutorial \
      --zone=europe-west4-a
    
    2018/04/28 16:16:23 WARNING: Setting zone to "europe-west4-a"
    No instances currently exist.
        Compute Engine VM:     --
        Cloud TPU:             --
    
  4. Execute gsutil conforme mostrado. Substitua bucket-name pelo nome do bucket do Cloud Storage criado para este tutorial:

    $ gsutil rm -r gs://bucket-name
    

A seguir

Neste tutorial, você treinou o modelo RetinaNet usando um conjunto de dados de amostra. Os resultados deste treinamento (na maioria dos casos) não podem ser usados para inferência. Para usar um modelo para inferência, é possível treinar os dados em um conjunto de dados disponível publicamente ou no seu próprio conjunto de dados. Os modelos treinados nas Cloud TPUs exigem que os conjuntos de dados estejam no formato TFRecord.

Use a amostra da ferramenta de conversão de conjunto de dados para converter um conjunto de dados de classificação de imagem no formato TFRecord. Se você não estiver usando um modelo de classificação de imagem, converta manualmente o conjunto de dados para o TFRecord. Para mais informações, consulte TFRecord e tf.Example

Ajuste de hiperparâmetros

Para melhorar o desempenho do modelo com o conjunto de dados, é possível ajustar os hiperparâmetros. Encontre informações sobre hiperparâmetros comuns a todos os modelos compatíveis com a TPU no GitHub. As informações sobre hiperparâmetros específicos do modelo podem ser encontradas no código-fonte de cada modelo. Para mais informações sobre o ajuste de hiperparâmetros, consulte Visão geral do ajuste de hiperparâmetros, Como usar o serviço de ajuste de hiperparâmetros e Ajustar hiperparâmetros.

Inferência

Depois de treinar o modelo, é possível usá-lo para inferência (também chamada de previsão). O AI Platform é uma solução baseada em nuvem para desenvolver, treinar e implantar modelos de machine learning. Depois que um modelo é implantado, use o serviço do AI Platform Prediction.

Treinar com diferentes tamanhos de imagem

Para explorar, use uma rede de backbone maior. Por exemplo, ResNet-101 em vez de ResNet-50. Uma imagem maior de entrada e um backbone mais potente produzirão um modelo mais lento, porém, mais preciso.

Usar uma base diferente

Se preferir, faça o pré-treinamento de um modelo do ResNet no próprio conjunto de dados e use-o como base para o modelo do RetinaNet. Com um pouco mais de trabalho, também é possível alternar uma outra rede de backbone em vez do ResNet. Por fim, caso tenha interesse em implementar modelos próprios de detecção de objetos, essa rede pode ser uma boa base para fazer mais experimentos.