Como processar imagens grandes com o Cloud TPU

Este documento abrange o uso do particionamento de espaço com Cloud TPUs e TPUEstimator para possibilitar o treinamento para imagens e vídeos muito grandes.

O recurso de particionamento de espaço do Cloud TPU permite dividir uma imagem através de vários chips TPU. O recurso de particionamento de espaço permite que você escalone modelos para serem executados com imagens de entrada grandes demais para caber na memória disponível em um único núcleo de um chip acelerador.

Esse recurso distribui ativações em vários núcleos de TPU, permitindo que você escalone seu modelo para usar 2, 4, 8 ou até 16 núcleos para treinamento. O escalonamento do número de processadores disponíveis para treinamento elimina a necessidade de reduzir a resolução dos dados da imagem, o que proporciona maior precisão e desempenho de treinamento mais eficiente.

É possível aplicar o particionamento de espaço a tarefas como:

  • Segmentação de imagens 3D de tomografia computadorizada
  • análise de conteúdo em vídeo
  • detecção de objetos para executar autonomamente sem diminuir a resolução dos dados de imagem.

Como ativar o particionamento de espaço com a TPUEstimator

O compilador otimizador de XLA para Tensorflow gerencia automaticamente as comunicações entre todos os núcleos do Cloud TPU. Não é necessária nenhuma alteração no código para ativar o particionamento de espaço em um modelo. Só é necessário configurar como particionar cada tensor de entrada na TPUConfig, já que a TPUEstimator é compatível com a API de particionamento de espaço.

imagem

Exemplo

O código a seguir mostra um exemplo de TPUConfig de particionamento de espaço em quatro níveis para um modelo de classificação de imagem. O tensor é dividido em quatro partes ao longo da dimensão de altura, supondo que o tensor tenha a forma (lote, altura, largura, canal).

tpu_config=tpu_config.TPUConfig(
    iterations_per_loop=100,
    num_cores_per_replica=4,
    per_host_input_for_training=tpu_config.InputPipelineConfig.PER_HOST_V2,
    input_partition_dims=[[1, 4, 1, 1], None]]

Para o particionamento de espaço, o pipeline de entrada precisa estar no formato tf.data e o valor per_host_input_for_training value (train_batch_size) precisa estar configurado para PER_HOST_V2. O num_cores_per_replica especificado determina o número máximo de partições em que é possível dividir uma imagem.

Os valores de input_partition_dims fornecem uma lista com dois elementos: feature_partition_dims e label_partition_dims, que descrevem como particionar os tensores de entrada. A estrutura de feature_partition_dims e label_partition_dims precisa corresponder à estrutura de recursos e rótulos de input_fn. Especifique "none" para a partição de rótulo para que os rótulos não sejam divididos.

Como executar modelos de referência com particionamento de espaço

Detecção de objetos em 2D

O RetinaNet é um modelo de detecção de objetos que localiza objetos em imagens com uma caixa delimitadora e classifica os objetos identificados. Treinar o modelo com imagens muito grandes é um desafio, porque a maior resolução que cabe em um único núcleo do Cloud TPU (com lote de 8 por dispositivo) é 1280 x 1280. É possível treinar imagens quatro vezes maiores, particionando espacialmente o modelo em oito núcleos de TPU, como mostrado.

Tamanho da imagem Tipo de TPU Núcleos de TPU Tamanho global do lote Tempo de duração
1280x1280 v3-8 8 64 910 ms
2560x2560 v3-64 64 64 822 ms

Segmentação de imagens em 3D

3D UNet é um modelo conhecido de segmentação 3D densa que tem sido usado na área médica. A resolução original de uma imagem de tomografia computadorizada pode chegar a 256x256x256, o que não cabe em um único núcleo TPU. Por isso, os pesquisadores geralmente precisam diminuir a escala das imagens. É possível encaixar imagens com a resolução original usando o particionamento de espaço em 16 níveis da TPU.

Tamanho da imagem Tipo de TPU Núcleos de TPU Tamanho global do lote Tempo de duração
128x128x128 v3-8 8 32 3.428s
2560x2560 v3-64 64 32 3.02s

Outros recursos