t Ejecuta TensorFlow en Cloud TPU

En este documento, se describen los detalles sobre la ejecución de modelos de TensorFlow en Cloud TPU. La forma de entrenar los modelos de TensorFlow en Cloud TPU depende de la arquitectura de Cloud TPU que uses. Para obtener más información, consulta Arquitectura del sistema.

Descripción general

Cuando ejecutas tu aplicación, TensorFlow genera un gráfico de computación y lo envía a un nodo TPU mediante gRPC. El nodo TPU compila el grafo de cálculo justo a tiempo y envía el objeto binario del programa a uno o más dispositivos de TPU para su ejecución. Las entradas del modelo se suelen almacenar en Cloud Storage. El nodo TPU transmite los datos de entrada a uno o más dispositivos de TPU para el entrenamiento.

En el diagrama de bloques a continuación, se muestra la arquitectura de software de Cloud TPU que consiste en el modelo de red neuronal, el cliente de TensorFlow, el servidor de TensorFlow y el compilador XLA.

image

API de TensorFlow de Cloud TPU

Keras es una API de aprendizaje profundo de alto nivel escrita en Python que se ejecuta sobre la plataforma de aprendizaje automático de TensorFlow. Cuando escribas un modelo de red neuronal que utilice Cloud TPU, recomendamos usar Keras.

Estimador de TPU (heredado)

Los estimadores de TPU son un conjunto de API heredadas de alto nivel creadas en estimadores que simplifican la compilación de modelos para Cloud TPU y que obtienen el máximo rendimiento de la TPU. Si bien el estimador de TPU aún es compatible con TensorFlow 1.x, recomendamos usar Keras en su lugar.

Servidor de TensorFlow

Un servidor de TensorFlow se ejecuta en un servidor de Cloud TPU. Cuando el servidor recibe un grafo de procesamiento del cliente de TensorFlow, realiza estas acciones:

  1. Carga entradas desde Cloud Storage.
  2. Divide el grafo en partes que se pueden ejecutar en una Cloud TPU y otras que se deben ejecutar en una CPU.
  3. Genera operaciones de XLA que corresponden al subgrafo que se ejecutará en Cloud TPU.
  4. Invoca el compilador XLA.

Para obtener más información sobre cómo entrenar un modelo de TensorFlow en nodos TPU, consulta la guía de inicio rápido de Cloud TPU.

VM de Cloud TPU

Con las VM de Cloud TPU, no hay VM del usuario y tu modelo y código se ejecutan en una VM en la máquina anfitrión de TPU. Puedes establecer una conexión SSH directamente en el host de TPU. Puedes crear tu cliente y servicio de TensorFlow en la misma VM o en VM diferentes. Puedes ejecutar código arbitrario en el host de TPU para realizar un procesamiento previo, etcétera.

image

Para obtener más información sobre cómo entrenar un modelo de TensorFlow en VM de TPU, consulta la guía de inicio rápido de Tensorflow sobre VM de Cloud TPU.