Procesa imágenes grandes con Cloud TPU

En este documento se explica el uso de la partición espacial con Cloud TPU y TPUEstimator para habilitar el entrenamiento de imágenes y videos muy grandes.

La función de partición espacial de Cloud TPU permite dividir una sola imagen en varios chips TPU. La partición espacial te permite escalar con facilidad los modelos a fin de ejecutarlos con imágenes de entrada que son demasiado grandes para la memoria disponible de un solo núcleo en un chip de aceleración.

Esta función distribuye activaciones en varios núcleos de TPU, lo que te permite escalar el modelo a fin de usar 2, 4, 8 o 16 núcleos para el entrenamiento. Cuando se escala la cantidad de procesadores disponibles para el entrenamiento, no es necesario realizar una reducción de muestreo de los datos de imágenes, lo que proporciona una mayor precisión y un rendimiento de entrenamiento más eficiente.

Puedes aplicar la partición espacial a tareas como las siguientes:

  • Segmentación de imágenes de escaneos de tomografía computarizada (TC) en 3D
  • Análisis de contenido de video
  • Detección de objetos para la conducción autónoma sin realizar una reducción de muestreo de los datos de imágenes

Habilita la partición espacial con TPUEstimator

En TensorFlow, el compilador de optimización XLA controla de forma automática las comunicaciones entre todos los núcleos de Cloud TPU. No es necesario cambiar el código para habilitar la partición espacial de un modelo. Debido a que TPUEstimator admite la API de partición espacial, solo debes configurar el modo de particionar cada tensor de entrada en TPUConfig.

image

Ejemplo

En el siguiente código, se muestra un ejemplo de TPUConfig de la partición espacial en cuatro partes para un modelo de clasificación de imágenes. El tensor se divide en cuatro partes a lo largo de la dimensión de altura (se supone que el tensor tiene la forma [lote, altura, ancho, 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 la partición espacial, la canalización de entrada debe estar en formato tf.data y el valor per_host_input_for_training (el train_batch_size) debe configurarse como PER_HOST_V2. El valor num_cores_per_replica que especifiques determina la cantidad máxima de particiones en las que puedes dividir una imagen.

Los valores input_partition_dims proporcionan una lista con dos elementos: feature_partition_dims y label_partition_dims que describen cómo particionar los tensores de entrada. La estructura de feature_partition_dims y label_partition_dims debe coincidir con la estructura de características y etiquetas de input_fn. Especifica “None” en la partición de etiquetas para que estas no se dividan.

Ejecuta modelos de referencia con partición espacial

Detección de objetos 2D

RetinaNet es un modelo de detección de objetos que localiza objetos en imágenes con un cuadro de límite y clasifica los objetos identificados. Entrenar el modelo con imágenes muy grandes es un desafío, ya que la imagen más grande que puede caber en una Cloud TPU de núcleo único (con 8 lotes por dispositivo) es de 1280x1280. Puedes entrenar imágenes que sean 4 veces más grandes si particionas de forma espacial el modelo en 8 núcleos de TPU, como se muestra.

Tamaño de la imagen Tipo de TPU Núcleos de TPU Tamaño del lote global Duración del paso
1280x1280 v3-8 8 64 910 ms
2560x2560 v3-64 64 64 822 ms

Segmentación de imágenes 3D

3D UNet es un popular modelo de segmentación 3D denso muy usado en el ámbito de las imágenes médicas. La resolución original de las imágenes de TC puede ser hasta de 256x256x256, lo que no cabe en un núcleo de TPU único. Por lo tanto, los investigadores deben hacer una reducción de muestreo en las imágenes. Con la partición espacial de TPU, puedes adaptar directamente las imágenes de resolución original mediante la partición espacial en 16 partes.

Tamaño de la imagen Tipo de TPU Núcleos de TPU Tamaño del lote global Duración del paso
128x128x128 v3-8 8 32 3.428 s
2560x2560 v3-64 64 32 3.02 s

Recursos adicionales