Traiter des images volumineuses avec Cloud TPU

Ce document présente l'utilisation du partitionnement spatial avec Cloud TPU et TPUEstimator pour permettre l'entraînement des images et des vidéos volumineuses.

La fonctionnalité de partitionnement spatial de Cloud TPU permet de répartir une image unique sur plusieurs puces TPU après fractionnement. Le partitionnement spatial facilite la mise à l'échelle des modèles à exécuter avec des images d’entrée trop volumineuses pour la mémoire disponible d'un seul cœur sur une puce d’accélération matérielle.

Cette fonctionnalité répartit les activations sur plusieurs cœurs TPU, ce qui vous permet d'adapter votre modèle pour utiliser 2, 4, 8 ou même 16 cœurs pour l'entraînement. L'augmentation du nombre de processeurs disponibles rend inutile le sous-échantillonnage des données d'image, ce qui améliore la précision et les performances de l'entraînement.

Le partitionnement spatial peut être appliqué à des tâches telles que :

  • la segmentation d'images 3D par tomodensitométrie (CT) ;
  • l'analyse de contenu vidéo ;
  • la détection d'objets pour la conduite autonome sans avoir à sous-échantillonner vos données d'image.

Activer le partitionnement spatial avec TPUEstimator

Dans TensorFlow, le compilateur d'optimisation XLA pour Tensorflow gère automatiquement les communications entre tous les cœurs Cloud TPU. L'activation du partitionnement spatial d'un modèle ne nécessite aucune modification du code. TPUEstimator est compatible avec l'API de partitionnement spatial. Il suffit de configurer chaque Tensor d'entrée pour son partitionnement dans TPUConfig.

image

Exemple

Le code suivant présente un exemple TPUConfig de partitionnement spatial quadridirectionnel pour un modèle de classification d'image. Le Tensor est divisé en quatre parties le long de la dimension de hauteur (en partant du principe que le Tensor détient la forme [lot, hauteur, largeur, 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]]

Pour le partitionnement spatial, le pipeline d'entrée doit être au format tf.data et la valeur per_host_input_for_training (le paramètre train_batch_size) doit être définie sur PER_HOST_V2. La valeur num_cores_per_replica que vous spécifiez détermine le nombre maximal de partitions dans lesquelles vous pouvez fractionner une image.

Les valeurs input_partition_dims fournissent une liste associée à deux éléments : feature_partition_dims et label_partition_dims qui décrivent comment partitionner les Tensors d'entrée. La structure de feature_partition_dims et label_partition_dims doit correspondre à la structure des fonctionnalités et des étiquettes de input_fn. Indiquez "None" pour la partition des étiquettes afin que celles-ci ne soient pas fractionnées.

Exécuter des modèles de référence avec partitionnement spatial

Détection d'objets 2D

RetinaNet est un modèle de détection d'objets qui localise les objets dans des images au moyen d'un cadre de délimitation et classe les objets identifiés. L'entraînement du modèle avec des images très volumineuses est un défi, car la taille maximale d'image pouvant être traitée sur un seul cœur Cloud TPU (8 par appareil pour un traitement par lot) est de 1 280 x 1 280. Il est possible d'entraîner des images quatre fois plus grandes en effectuant un partitionnement spatial du modèle sur 8 cœurs TPU, comme illustré.

Taille d'image Type de TPU Cœurs TPU Taille du lot global Durée des pas
1280 x 1280 v3-8 8 64 910 ms
2 560 x 2 560 v3-64 64 64 822 ms

Segmentation d'image 3D

3D UNet est un modèle de segmentation 3D dense couramment utilisé dans le domaine de l'imagerie médicale. La résolution d'origine des images CT pouvant atteindre 256 x 256 x 256, elles sont impossibles à traiter sur un seul cœur TPU. Les chercheurs doivent donc souvent sous-échantillonner les images. Le partitionnement spatial TPU vous permet d'ajuster directement les images dans la résolution d'origine à l'aide d'un partitionnement spatial à 16 voies.

Taille d'image Type de TPU Cœurs TPU Taille du lot global Durée des pas
128 x 128 x 128 v3-8 8 32 3,428 s
2 560 x 2 560 v3-64 64 32 3,02 s

Autres ressources