t Exécuter TensorFlow sur Cloud TPU

Ce document décrit en détail l'exécution de modèles TensorFlow sur Cloud TPU. La façon dont vous entraînez les modèles TensorFlow sur Cloud TPU dépend de l'architecture Cloud TPU que vous utilisez. Pour plus d'informations, 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 à 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. Le nœud TPU transmet les données d'entrée à un ou plusieurs appareils TPU pour l'entraînement.

Le diagramme ci-dessous illustre l'architecture logicielle Cloud TPU constituée du 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 estimateurs 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 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. Générer des opérations XLA correspondant au sous-graphe à exécuter sur Cloud TPU
  4. Il appelle le compilateur XLA.

Pour plus d'informations 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

Il n'y a aucune 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 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 le prétraitement, etc.

image

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