Cómo utilizar las TPU de Google Cloud en tus proyectos de Inteligencia Artificial y un poco de la historia del hardware a medida de Google

En 1996, Sergey Brin y Larry Page construyeron su primer prototipo para probar una idea un poco “loca” ;) que tenía Larry para indexar las páginas web. En aquella época no era fácil ni barato montar un servidor de las características que ellos requerían, así que decidieron construirlo ellos mismos. Necesitaban por lo menos 10 discos duros de 4GB (si GB, no Teras, estos eran los más grandes de la época) pero no había cajas disponibles para esta configuración tan específica. Así que decidieron utilizar … piezas de LEGO. Y parece que funcionó, ya que en 1998 fundaron una empresa que a lo mejor te suena, se llamaba Google. 

Figura 1: Cómo utilizar las TPU de Google Cloud en tus proyectos de

Inteligencia Artificial y un poco de la historia del hardware a medida de Google

Ahora que ya tenían una compañía era el momento de minimizar costes para su primera infraestructura así que construyeron 30 racks con las partes más baratas que encontraron en el mercado utilizando corcho (por eso se le quedó en el argot el nombre de servidor “corkboard” de Google) para aislar los diferentes módulos. 

Figura 2: Izquierda, primer prototipo con la caja de LEGO

y a la derecha, el primer servidor en producción de Google.

Este rack tenía hardware de todo tipo como Sun Microystems Ultra II, Pentium II, IBM RS/6000, discos de 9GB y 5GB, etcétera. Como podéis observar, arquitecturas y configuraciones varias. Y así comenzó la “tradición” de Google de crear su propia configuración de hardware y software para construir sus propios servidores. Esto continua incluso hoy día con la arquitectura de las TPU o las Tensor Processing Units de las cuales hablaremos a continuación.

Figura 3: Libro de "Microhistorias: anécdotas y curiosiades de la historia
de la informática (y los hackers)" de Fran Ramírez y Rafel Troncoso 0xWord.

Pero antes, si te ha gustado esta introducción retro-informática ;) recuerda que tenemos el libro "Microhistorias: Anécdotas y Curiosidades de la Informática ( y los hackers)" de 0xWord donde contamos historias como esta de hackers y personalidades de la Informática.

Tensor Processing Units

Con el auge de la Inteligencia Artificial (y en concreto del Machine Learning), de la que Google es posiblemente uno de los mayores líderes mundiales en esta tecnología, vieron la necesidad de construir también su propio Hardware Custom-Made (término anglosajón que significa “a medida”). Así que en 2017, Google anunció las TPU o Tensor Processing Units. 

Figura 4: TPU v3, 4 chips, 2 cores por cada chip.

Pero esta vez no era un servidor hecho de piezas de otros fabricantes, ahora habían construido su propio chip, un ASIC (Custom Application-Specific Integrated Circuit) enfocado a las matemáticas del Machine Learning. Y más específicamente, está diseñado para trabajar directamente con una librería matemática creada también por Google llamada Tensorflow.

CPU vs GPU

Una CPU es un hardware diseñado para un uso general de programación y control (componentes del sistema como gráficos, memoria, accesos a periféricos, etc) que suele tener varios cores o núcleos (por ejemplo, 64 núcleos tiene el procesador más potente en 2002 que puedes comprar, el AMD Ryzen Threadripper 3995WX) capaz de realizar todo tipo de operaciones normalmente secuenciales. Una GPU es un procesador específico para acelerar los procesos gráficos de la CPU utilizando pequeños cores (aquí hablamos de miles) trabajando en paralelo (Parallel Processing) realizando tareas distintas. 


Figura 5: NVIDIA: Adam and Jamie explain parallel processing on GPU's

Estos cores están especializados en realizar operaciones matemáticas, y más en concreto operaciones con matrices, las cuales son perfectas para por ejemplo Deep Learning y en general para la Inteligencia Artificial. En concreto, hablamos de la manipulación de escalares, vectores y matrices además de sus operaciones asociadas. En este enlace de NVIDIA explican a fondo la diferencia entre CPU y GPU. Pero la mejor explicación que podemos encontrar nos la ofrecen los Cazadores de Mitos en el vídeo de la Figura 5.

Tensores y las TPU

 
Un "tensor" es un objeto matemático que nos permite mapear y manipular elementos como escalares, vectores, etcétera. En otras palabras, un tensor es una matriz de 0 dimensiones (un número/escalar, por ejemplo), 1 dimensión (vector), 2 dimensiones (matriz), etcétera. Las dimensiones se denominan rango. Si utilizas Python y conoces Numpy, un tensor es muy parecido a una matriz de n-dimensiones de esta librería de Python. En este enlace puedes encontrar más información sobre qué es un tensor.

Figura 6: Representación gráfica de Tensores.

Durante el proceso de entrenamiento de una red neuronal usamos continuamente operaciones de multiplicación y suma de pesos (valores de las conexiones de las diferentes neuronas) así como por ejemplo durante las funciones de activación. Y todas ellas podemos efectuarlas utilizando matrices.

Figura 7: Unificación de los cálculos de una red neuronal en una matriz.

Y aquí es donde aparecen las TPU para revolucionar el mundo del Machine Leaerning, un hardware específico creado por Google para realizar únicamente operaciones para redes neuronales (es decir, potenciar la fase de entrenamiento principalmente) a unas velocidades realmente increíbles (hablamos de pasar de meses de entrenamiento a días). 

Figura 8: Libro de Machine Learning aplicado a Ciberseguridad de
Carmen Torrano, Fran Ramírez, Paloma Recuero, José Torres y Santiago Hernández

Además, las TPU pueden realizar computaciones de álgebra lineal, algo que las GPUs no pueden realizar al estar enfocadas en operaciones con matrices. En general, las TPU incluye los siguientes elementos:

• MXU o Matrix Multiplier Unit: Unidad para todo tipo de cálculos con matrices.

• UB o Unified Buffer: Registros de 24MB de SRAM (misma funcionalidad que los que conocemos para las CPU).

• AU o Activation Unit: Funciones de activación, pero por hardware lo que aumenta exponencialmente su rendimiento.

Hoy en día tenemos la serie de TPUv1i a TPUv4i dedicadas a la inferencia en redes neuronales profundas (inferencia es el proceso de utilizar DNN o Redes Neuronales Profundas para realizar predicciones con datos que no conoce) y por otro lado tenemos la serie TPUv1 a TPUv4 dedicadas al entrenamiento de las DNN. Y como antes se ha mencionado, este potencial es aún mayor al combinar el framework Tensorflow con las TPU. 

Figura 9: Esquema interno de una TPU v3.

Mencionar que, de momento, las TPU sólo se pueden acceder desde el Cloud de Google, y más concretamente desde Google Colab, como veremos a continuación. Este es el enlace principal de Google Cloud y el servicio de TPUs.

¿Cómo podemos usar las TPU en nuestros proyectos?

En Google Colab es posible utilizar una TPU de 8 núcleos manera gratuita. Esta se puede seleccionar desde la pestaña "Entorno de ejecución –> Cambiar tipo de entorno de ejecución –> Acelerador por hardware –> TPU".

Figura 10: Seleccionar TPU en Google Colab.

Sin embargo, al contrario que con las GPUs, es necesario hacer unos pasos más para poder utilizarlas, a pesar de haber seleccionado ya el entorno de ejecución. En primer lugar, debes conectarte a un clúster TPU:

import tensorflow as tf
tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
print("Dispositivo:", tpu.master())
  

Ahora, la ejecución de tu notebook estará conectada a una TPU. En este momento es necesario crear un scope strategy (TPUStrategy), para que el modelo se replique una vez por cada núcleo de la TPU, manteniendo sincronizadas las réplicas:

strategy = tf.distribute.TPUStrategy(tpu) 
print(f"Número de replicas: {strategy.num_replicas_in_sync}")

En este momento es importante definir el modelo que queramos y compilarlo abriendo un scope:

with strategy.scope():

	inputs = keras.Input(…)
	…
	outputs =  layers.Dense(…)
	model = keras.Model(inputs, outputs)

	model.compile(…)

A continuación, ya podemos comenzar el proceso de aprendizaje del modelo (esta línea ya se ejecutará fuera del scope creado en el paso anterior). A la hora de entrenar en TPU, cuanto más grande sea el tamaño del batch, más se aprovecharán los núcleos de ésta, por lo que es frecuente encontrarse batches de gran tamaño:

model.fit(x_train, y_train, batch_size=1024,…) 

Es muy importante tener en cuenta que cuando se utilizan TPU en Google Colab, se configuran dos máquinas virtuales. Una de ellas aloja la ejecución de tu notebook, mientras que la otra es donde vive la TPU. Es por esto por lo que no podremos entrenar nuestra red neuronal como solemos hacer cargando los datos del disco local, ya que desde la máquina virtual donde está la TPU no se puede acceder al disco de la máquina que aloja el notebook.

Existen dos opciones, o bien el dataset es lo suficientemente pequeño como para poder alojarlo en la memoria de la máquina virtual en un array (este es el ejemplo que se ha mostrado en este artículo, teniendo lo datos en los arrays (x_train, y_train)), o la otra opción consiste en alojar los datos en un bucket de Google Cloud Storage (GCS) y crear un dataset que vaya leyendo los datos de este lugar, sin descargarlos en ningún momento de manera local, ya que el entorno de ejecución TPU está configurado para poder leer datos de GCS.

Por otra parte, si tu modelo es relativamente pequeño, no estará aprovechando al máximo recursos de la TPU a no ser que establezcas tamaños de batch bastante grandes (de más de 10.000 muestras por batch). En este caso, como se realizarán menos actualizaciones de pesos que con un tamaño de batch estándar, ya que se procesan muchas más muestras antes de actualizarse estos, sería necesario elevar el paso de aprendizaje, porque se harán actualizaciones más realistas y podemos por lo tanto hacer pasos más grandes (como se habrán visto muchas más muestras, se “adivinará” mejor la minimización del error).

Sin embargo, existe otra alternativa mejor para no tener que utilizar batches tan grandes, que consiste en el llamado Step Fusing. La idea es ejecutar múltiples pasos de entrenamiento durante cada paso de ejecución de la TPU. Básicamente, hacer más trabajo entre dos viajes de ida y vuelta desde la memoria de la máquina virtual a la TPU. Para ello, cuando compilemos el modelo, es necesario añadir el argumento steps_per_execution:

model.compile(..., steps_per_execution=N)

Por ejemplo, steps_per_execution=6 ejecutará seis pasos de entrenamiento durante cada ejecución de la TPU.

Conclusiones

Las TPU son un dispositivo innovador el cual seguro que dará lugar a otros específicos orientados a la ejecución de modelos de Machine Learning, más allá de las clásicas GPU. La velocidad tanto de procesamiento como de entrenamiento en la Inteligencia Artificial son dos factores fundamentales que marcan la posible aplicabilidad o no de una solución que resuelva un problema de la sociedad.  Por este motivo son tan importantes este tipo de avances tecnológicos. Y si quieres aprender cómo la Inteligencia Artificial se puede aplicar en el mundo de la Ciberseguridad, te recomendamos este libro de 0xWord: Machine Learning Aplicado a Ciberseguridad.

Saludos,

Autores:

Fran Ramírez, es investigador de seguridad y miembro del equipo de Ideas Locas en CDO en Telefónica, co-autor del libro "Microhistorias: Anécdotas y Curiosidades de la historia de la informática (y los hackers)", del libro "Docker: SecDevOps", también de "Machine Learning aplicado a la Ciberseguridad” además del blog CyberHades. Puedes contactar con Fran Ramirez en MyPublicInbox.

 Contactar con Fran Ramírez en MyPublicInbox

Javier del Pino Díaz (Intership en Ideas Locas CDO de Telefónica)

Contactar con Javier del Pino Díaz en MyPublicInbox