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 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.

Cuando ejecutas tu aplicación, TensorFlow genera un grafo de procesamiento y lo envía a una TPU. El host de TPU compila el grafo de cómputo 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 host de TPU transmite los datos de entrada a uno o más dispositivos de TPU para el entrenamiento.

VM de Cloud TPU

Con la arquitectura de 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 al 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

Si deseas obtener más información para entrenar un modelo de TensorFlow en las VM de TPU, consulta la Guía de inicio rápido sobre Tensorflow de VM de Cloud TPU.

Nodos de Cloud TPU

Con la arquitectura del nodo de Cloud TPU, un host de TPU consta de una VM de Compute Engine y un nodo TPU. La VM de Compute Engine ejecuta un cliente de TensorFlow y tu modelo. La VM se comunica con el nodo TPU a través de gRPC. El nodo TPU ejecuta el servidor de TensorFlow y el compilador XLA.

En el siguiente diagrama de bloques, se muestra la arquitectura de software del nodo 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 Cloud TPU para TensorFlow

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, te recomendamos usar Keras.

Estimador de TPU (heredado)

Los estimadores de TPU son un conjunto de API de alto nivel heredadas 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 de parte 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.

Si deseas obtener más información para entrenar un modelo de TensorFlow en los nodos TPU, consulta la Guía de inicio rápido de Cloud TPU.