Exécuter TensorFlow sur Cloud TPU

Ce document décrit en détail comment exécuter des modèles TensorFlow sur des Cloud TPU. La manière dont vous entraînez les modèles TensorFlow sur Cloud TPU dépend de l'architecture Cloud TPU que vous utilisez. Pour en savoir plus, consultez la page Architecture du système.

Lorsque vous exécutez votre application, TensorFlow génère un graphe de calcul et l'envoie à un TPU. L'hôte TPU compile le graphe de calcul à la volée et envoie le binaire du programme à un ou plusieurs appareils TPU pour exécution. Les entrées du modèle sont souvent stockées dans Cloud Storage. L'hôte TPU transmet les données d'entrée à un ou plusieurs appareils TPU pour l'entraînement.

VM Cloud TPU

Avec l'architecture de VM Cloud TPU, il n'y a pas de VM utilisateur, et votre modèle et votre code s'exécutent dans une VM sur la machine hôte TPU. Vous pouvez vous connecter directement à l'hôte TPU en SSH. Vous pouvez créer votre client et votre service TensorFlow sur la même VM ou sur des VM différentes. Vous pouvez exécuter du code arbitraire sur l'hôte TPU pour effectuer un prétraitement, etc.

image

Pour en savoir plus sur l'entraînement d'un modèle TensorFlow sur des VM TPU, consultez le guide de démarrage rapide de TensorFlow pour les VM Cloud TPU.

Nœuds Cloud TPU

Avec l'architecture de nœud Cloud TPU, un hôte TPU se compose d'une VM Compute Engine et d'un nœud TPU. La VM Compute Engine exécute un client TensorFlow et votre modèle. La VM communique avec le nœud TPU via gRPC. Le nœud TPU exécute le serveur TensorFlow et le compilateur XLA.

Le schéma de blocs suivant illustre l'architecture logicielle Cloud TPU à base de modèle de réseau de neurones, du client TensorFlow, du serveur TensorFlow et du compilateur XLA.

image

API Cloud TPU TensorFlow

Keras est une API de deep learning de haut niveau écrite en Python, qui s'exécute sur la plate-forme de machine learning TensorFlow. Lorsque vous écrivez un modèle de réseau de neurones utilisant Cloud TPU, nous vous recommandons d'utiliser Keras.

TPU Estimator (ancien)

Les TPU Estimators sont un ensemble d'anciennes API de haut niveau qui s'appuient sur des Estimators simplifiant la création de modèles pour Cloud TPU et tirant parti des performances maximales des TPU. Bien que TPU Estimator soit toujours compatible avec Tensorflow 1.x, nous vous recommandons plutôt d'utiliser Keras.

Serveur TensorFlow

Un serveur TensorFlow s'exécute sur un serveur Cloud TPU. Lorsque le serveur reçoit un graphe de calcul du client TensorFlow, il réalise les actions suivantes :

  1. Il charge les entrées depuis Cloud Storage.
  2. Il partitionne le graphe en sections pouvant s'exécuter sur un Cloud TPU et en sections devant s'exécuter sur un processeur.
  3. Il génère les opérations XLA correspondant au sous-graphe à exécuter sur Cloud TPU.
  4. Il appelle le compilateur XLA.

Pour en savoir plus sur l'entraînement d'un modèle TensorFlow sur des nœuds TPU, consultez le guide de démarrage rapide de Cloud TPU.