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.

Présentation

Lorsque vous exécutez votre application, TensorFlow génère un graphe de calcul et l'envoie à un nœud TPU via gRPC. Le nœud TPU compile le graphe de calcul au plus vite 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. Le nœud TPU transmet les données d'entrée à un ou plusieurs appareils TPU pour l'entraînement.

Le diagramme de bloc ci-dessous illustre l'architecture logicielle Cloud TPU, qui comprend le modèle de réseau de neurones, le client TensorFlow, le serveur TensorFlow et le 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 qui utilise Cloud TPU, nous vous recommandons d'utiliser Keras.

Estimator de TPU (ancien)

Les TPU Estimators sont un ensemble d'anciennes API de haut niveau qui s'appuient sur des Estimators pour simplifier la création de modèles pour Cloud TPU et extraire les performances maximales des TPU. L'Estimator de TPU est toujours compatible avec TensorFlow 1.x, mais nous vous recommandons d'utiliser Keras à la place.

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 effectue les actions suivantes:

  1. Il charge les entrées depuis Cloud Storage.
  2. Partitionner le graphe en parties pouvant s'exécuter sur un Cloud TPU et en parties devant être exécutées 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.

VM Cloud TPU

Avec les VM Cloud TPU, il n'y a pas de VM utilisateur, et votre modèle et votre code s'exécutent sur une VM sur la machine hôte TPU. Vous pouvez vous connecter en SSH directement à l'hôte TPU. 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.