Como converter um conjunto de dados de classificação de imagem para uso com o Cloud TPU

Neste tutorial, descrevemos como usar o script de amostra de conversor de dados de classificação de imagens para converter um conjunto de dados brutos de classificação de imagens no formato TFRecord usado para treinar modelos do Cloud TPU.

Os TFRecords tornam a leitura de arquivos grandes do Google Cloud Storage mais eficiente do que ler cada imagem como um arquivo individual. É possível usar o TFRecord sempre que você usa um pipeline tf.data.Dataset.

Consulte os seguintes documentos do TensorFlow para mais informações sobre como usar o TFRecord:

Se você usa o framework PyTorch ou JAX e não usa o Google Cloud Storage para armazenamento de conjunto de dados, talvez não tenha a mesma vantagem dos TFRecords.

Visão geral da conversão

A pasta de classificação de imagens no repositório do conversor de dados no GitHub contém o script converter, image_classification_data.py e uma implementação de amostra, simple_example.py, que podem ser copiados e modificados para você fazer a própria conversão de dados.

A amostra de conversão de dados de classificação de imagem define duas classes, ImageClassificationConfig e ImageClassificationBuilder. Essas classes são definidas em tpu/tools/data_converter/image_classification_data.py.

ImageClassificationConfig é uma classe base abstrata. Você usa a subclasse ImageClassificationConfig para definir a configuração necessária para criar uma instância de ImageClassificationBuilder.

ImageClassificationBuilder é um builder de conjunto de dados do Tensorflow para conjuntos de dados de classificação de imagens. É uma subclasse de tdfs.core.GeneratorBasedBuilder. Ele recupera exemplos de dados do conjunto de dados e os converte em TFRecords. Os TFRecords são gravados em um caminho especificado pelo parâmetro data_dir para o método __init__ de ImageClassificationBuilder.

Em simples_example.py, SimpleDatasetConfig usa a subclasse ImageClassificationConfig, implementando propriedades que definem os modos compatíveis, o número de classes de imagem e um gerador de exemplo que produz um dicionário contendo dados de imagem e uma classe de imagem para cada exemplo no conjunto de dados.

A função main() cria um conjunto de dados de imagem gerados aleatoriamente e instancia um objeto SimpleDatasetConfig, especificando o número de classes e o caminho para o conjunto de dados no disco. Em seguida, main() instancia um objeto ImageClassificationBuilder, inserindo a instância SimpleDatasetConfig. Por fim, main() chama download_and_prepare(). Quando esse método é chamado, a instância ImageClassificationBuilder usa o gerador de exemplo de dados implementado por SimpleDatasetConfig para carregar cada exemplo e os salva em uma série de arquivos TFRecord.

Para uma explicação mais detalhada, consulte o Notebook do conversor de classificação.

Antes de começar

Como modificar a amostra de conversão de dados para carregar seu conjunto de dados

Para converter seu conjunto de dados no formato TFRecord, subclassifique a classe ImageClassificationConfig que define as seguintes propriedades:

  • num_labels: retorna o número de classes de imagem
  • supported_modes: retorna uma lista de modos compatíveis com seu conjunto de dados (por exemplo: teste, treinamento e validação)
  • text_label_map: retorna um dicionário que modela o mapeamento entre um rótulo de classe de texto e um rótulo de classe de número inteiro (SimpleDatasetConfig não usa essa propriedade, porque não requer um mapeamento)
  • download_path: o caminho do qual fazer o download do conjunto de dados (SimpleDatasetConfig não usa essa propriedade, o example_generator carrega os dados do disco)

Implemente a função do gerador example_generator. Esse método precisa produzir um dicionário contendo os dados da imagem e o nome da classe da imagem para cada exemplo. ImageClassificationBuilder usa a função example_generator() para recuperar cada exemplo e grava-os no disco no formato TFRecord.

Como executar a amostra de conversão de dados

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

    gsutil mb -p ${PROJECT_ID} -c standard -l us-central1 gs://bucket-name
    
  2. Inicie uma VM do Compute Engine usando o comando gcloud.

    $ gcloud compute tpus execution-groups create \
     --vm-only \
     --zone=us-central1-b \
     --name=imageclassificationconverter \
     --tf-version=2.5.0
    
    gcloud compute ssh imageclassificationconverter --zone=us-central1-b 

    Daqui em diante, o prefixo (vm)$ significa que é preciso executar o comando na instância de VM do Compute Engine.

  3. Instalar os pacotes necessários

    (vm)$ pip3 install opencv-python-headless pillow
    
  4. Crie as seguintes variáveis de ambiente usadas pelo script.

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    
    (vm)$ export CONVERTED_DIR=$HOME/tfrecords
    (vm)$ export GENERATED_DATA=$HOME/data
    (vm)$ export GCS_CONVERTED=$STORAGE_BUCKET/data_converter/image_classification/tfrecords
    (vm)$ export GCS_RAW=$STORAGE_BUCKET/image_classification/raw
    (vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/tpu/models"
    
  5. Altere para o diretório data_converter:

    (vm)$ cd /usr/share/tpu/tools/data_converter
    

Como executar o conversão de dados em um conjunto de dados falso

O script simple_example.py está localizado na pasta image_classification da amostra do conversor de dados. Executar o script com os parâmetros a seguir gera um conjunto de imagens falsas e as converte em TFRecords.

(vm)$ python3 image_classification/simple_example.py \
  --num_classes=1000 \
  --data_path=$GENERATED_DATA \
  --generate=True \
  --num_examples_per_class_low=10 \
  --num_examples_per_class_high=11 \
  --save_dir=$CONVERTED_DIR

Como executar o conversão de dados em um dos nossos conjuntos de dados brutos

  1. Crie uma variável de ambiente para o local dos dados brutos.

    (vm)$ export GCS_RAW=gs://cloud-tpu-test-datasets/data_converter/raw_image_classification
    
  2. Execute o script simple_example.py:

    (vm)$ python3 image_classification/simple_example.py \
    --num_classes=1000 \
    --data_path=$GCS_RAW \
    --generate=False \
    --save_dir=$CONVERTED_DIR
    

O script simple_example.py usa os seguintes parâmetros:

  • num_classes se refere ao número de classes no conjunto de dados. Usamos 1.000 aqui para corresponder ao formato ImageNet
  • generate determina se os dados brutos serão ou não gerados.
  • data_path se refere ao caminho onde os dados devem ser gerados se generate=True ou o caminho onde os dados brutos são armazenados se generate=False.
  • num_examples_per_class_low e num_examples_per_class_high determinam quantos exemplos cada classe gerar. O script gera um número aleatório de exemplos nesse intervalo.
  • save_dir se refere ao local onde os TFRecords salvos são armazenados. Para treinar um modelo no Cloud TPU, os dados precisam ser armazenados no Cloud Storage. Isso pode ser feito no Cloud Storage ou na VM.

Como renomear e mover os TFRecords para o Cloud Storage

O exemplo a seguir usa os dados convertidos com o modelo ResNet.

  1. Renomeie os TFRecord para o mesmo formato que os TFRecords ImageNet:

    (vm)$ cd $CONVERTED_DIR/image_classification_builder/Simple/0.1.0/
    (vm)$ sudo apt install rename
    
    (vm)$ rename -v 's/image_classification_builder-(\w+)\.tfrecord/$1/g' *
  2. Copie os TFRecords para o Cloud Storage:

    (vm)$ gsutil -m cp train* $GCS_CONVERTED
    (vm)$ gsutil -m cp validation* $GCS_CONVERTED