Ejecuta TensorFlow en Cloud TPU

En este documento, se describen los detalles para ejecutar modelos de TensorFlow en Cloud TPU. La forma en que entrenas 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 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 siguiente diagrama de bloques, 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.

imagen

API de TensorFlow de Cloud TPU

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

Estimador de TPU (heredado)

Los Estimadores de TPU son un conjunto de API heredadas de alto nivel que se basan en Estimadores, lo que simplifica la compilación de modelos para Cloud TPU y que extrae el rendimiento máximo de 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, el servidor realiza las siguientes acciones:

  1. Carga entradas desde Cloud Storage.
  2. Divide el grafo en partes que se pueden ejecutar en una Cloud TPU y otras que deben ejecutarse 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 los 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 una VM de 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 el procesamiento previo, etcétera.

imagen

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