TPUs de Google: el imparable avance del hardware especializado para ML.

El mundo de la inteligencia artificial y en concreto del machine learning (ML) ha evolucionado rápidamente en los últimos años.

Las mejoras en el hardware que incluyen mayor procesamiento acompañado de una reducción sustancial en el coste, unido a nuevos algoritmos que aportan nuevas capacidades y las hacen más eficientes, conforman los ingredientes perfectos para este crecimiento.

Ciertos procesos, sin embargo, siguen requiriendo de una potencia de cálculo que en muchas situaciones es difícil de proporcionar incluso por los procesadores más modernos (incluyendo CPUs, GPUs e incluso FPGAs de última generación), haciendo que el tiempo hasta que el proceso termina se alargue a horas o incluso días.

La tipología de muchas de las operaciones que se realizan durante el entrenamiento de un modelo de ML, o la misma inferencia, es bastante concreto y por lo general son operaciones no muy habituales en otro tipo de procesos.

Nos encontramos por lo tanto ante la necesidad de hardware especializado que aporte mayor eficiencia en este tipo de operaciones y aquí es donde entran en juego las tensor processing unit (TPU).

Hace unas semanas en Paradigma nos hicimos con una placa Google Coral Dev Board, un dispositivo que incluye una de estas TPU. En este artículo vamos a ver cómo funciona esta tecnología.

Las TPUs

Los tensor processing unit (TPU) son un hardware diseñado específicamente para resolver cierto tipo de operaciones muy frecuentes en procesos de machine learning, incluyendo el entrenamiento de los modelos o incluso la propia inferencia.

Google anunció en 2016 las tensor processing unit (TPU) dentro del marco del evento de Google I/O, y prometía un rendimiento 15-30X mayor junto con una mejora del 30-80X en cuanto a rendimiento por vatio que las CPUs y GPUs existentes en ese momento.

Para entender la forma de conseguir esta mejora de rendimiento, y generalizando para no tener que entrar mucho en detalle, podríamos decir una GPU se basa en un conjunto de ‘cores’ programables que proporcionan un alto rendimiento en operaciones con números con alta precisión. Esta arquitectura proporciona flexibilidad a la hora de ejecutar diversas tareas (renderizado de imágenes 3D, tratamiento de físicas…).

Por su parte, una TPU se centra en operaciones vectoriales y operaciones sobre matrices con números con baja precisión. Esta arquitectura hace que sean altamente eficientes en la multiplicación de grandes matrices, operaciones muy presentes al trabajar en ciertos tipos de redes de neuronas.

Google Cloud puso sus TPUs a disposición pública dentro de su servicio gestionado de infraestructura Cloud TPU tras el lanzamiento, y a día de hoy sigue evolucionando su plataforma con actualizaciones sobre su hardware.

La tecnología TPU está, además, a disposición y de forma gratuita en su entorno Google Colab, por lo que tenemos al alcance de la mano comprobar su funcionamiento de forma sencilla.

¿Cómo funciona?

La base de la tecnología TPU es su potencial para trabajar con matrices, así que vamos a ver un ejemplo en el que podamos sacarle partido.

Hace poco nuestro compañero Juan Iglesias nos enseñó cómo reconocer la escritura mediante dos alternativas: perceptron y redes convolucionales (CNN). Usaremos estas últimas para ver su funcionamiento, ya que encaja con el tipo de operación que buscamos.

Convolución

Las matrices de convolución suelen estar presentes en modelos que trabajan con imágenes, por ejemplo, en clasificación o detección; y en software de tratamiento de imágenes, como aquí, aunque se extiende a muchos otros campos.

Como Juan nos indicaba en su post, una convolución se puede interpretar “como la aplicación de filtros de diversos tipos”.

Supongamos que tenemos una imagen en blanco y negro de 6 x 6 pixeles, que se puede traducir fácilmente a una matriz de 1s y 0s:

Bien, ahora queremos aplicar un filtro (otra matriz, se denomina kernel) para obtener una serie de características específicas de la matriz original. Usaremos la siguiente matriz para nuestro ejemplo:

Simplificando, debemos deslizar la matriz filtro sobre la matriz original, por lo que la aplicación de la convolución de las dos matrices (la aplicación del filtro) se aplicaría de la siguiente forma:

Paso 1:

Paso 2

Paso 3

Al aplicar todos los pasos tendremos:

La generalización desde el punto de vista matemático para la convolución:

Pensando en una situación normal, estaríamos trabajando con imágenes a color (es decir, tendríamos las 3 matrices uno por cada color RGB), con tamaños de imagen mucho más grandes y es posible que con filtros de mayor tamaño… el volumen de operaciones a realizar se dispara.

La convolución es sólo un ejemplo de los tipos de operaciones en los que un procesador con elementos específicos para el trabajo con matrices puede aprovechar todo su potencia

El código

Veamos cómo podemos traducir esto a código para poder hacer nuestras pruebas.

Vamos a definir el modelo que utilizaremos:

model = tf.keras.Sequential(
     [
       tf.keras.layers.Reshape(input_shape=(IMG_ROWS*IMG_COLS,), target_shape=(IMG_ROWS, IMG_COLS, 1), name="image"),

       tf.keras.layers.Conv2D(filters=12, kernel_size=3, padding='same', use_bias=False), # no bias necessary before batch norm
       tf.keras.layers.BatchNormalization(scale=False, center=True), # no batch norm scaling necessary before "relu"
       tf.keras.layers.Activation('relu'), # activation after batch norm

       tf.keras.layers.Conv2D(filters=24, kernel_size=6, padding='same', use_bias=False, strides=2),
       tf.keras.layers.BatchNormalization(scale=False, center=True),
       tf.keras.layers.Activation('relu'),

       tf.keras.layers.Conv2D(filters=64, kernel_size=6, padding='same', use_bias=False, strides=2),
       tf.keras.layers.BatchNormalization(scale=False, center=True),
       tf.keras.layers.Activation('relu'),

       tf.keras.layers.Flatten(),
       tf.keras.layers.Dense(200, use_bias=False),
       tf.keras.layers.BatchNormalization(scale=False, center=True),
       tf.keras.layers.Activation('relu'),
       tf.keras.layers.Dropout(0.4), # Dropout on dense layer only

       tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')
     ])

Vemos varias capas del tipo Conv2D, estas son las capas de convolución del tipo que hemos analizado anteriormente y podemos ver que la primera, por ejemplo, tiene el parámetro kernel_size=3, indicando que es una matriz de 3x3.

No vamos a entrar en detalle en la configuración, os remito al post de Juan para más detalle.

Como indicaba anteriormente, podemos utilizar el entorno Google Colab para nuestras pruebas sobre la tecnología TPU, únicamente es necesario indicarlo en el Runtime que vayamos a utilizar (Menú Runtime -> Change Runtime Type).

Para poder hacer uso necesitamos autenticarnos:

IS_COLAB_BACKEND = 'COLAB_GPU' in os.environ  # this is always set on Colab, the value is 0 or 1 depending on GPU presence
if IS_COLAB_BACKEND:
 from google.colab import auth
 # Authenticates the Colab machine and also the TPU using your
 # credentials so that they can access your private GCS buckets.
 auth.authenticate_user()

E incluir el siguiente código, que nos servirá tanto si estamos utilizando un Runtime con TPU, GPU o únicamente CPU:

# Detect hardware
try:
 tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection
except ValueError:
 tpu = None
 gpus = tf.config.experimental.list_logical_devices("GPU")

# Select appropriate distribution strategy
if tpu:
 resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
 tf.config.experimental_connect_to_cluster(resolver)
 tf.tpu.experimental.initialize_tpu_system(resolver)
 strategy = tf.distribute.experimental.TPUStrategy(resolver)
 print('Running on TPU ', tpu.cluster_spec().as_dict()['worker']) 
elif len(gpus) > 1:
 strategy = tf.distribute.MirroredStrategy([gpu.name for gpu in gpus])
 print('Running on multiple GPUs ', [gpu.name for gpu in gpus])
elif len(gpus) == 1:
 strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU
 print('Running on single GPU ', gpus[0].name)
else:
 strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU
 print('Running on CPU')
print("Number of accelerators: ", strategy.num_replicas_in_sync)

A través de la strategy obtenida podremos utilizar nuestro modelo:

...
with strategy.scope():
   model = make_model()

Aquí os dejo el enlace al notebook (para poder importarlo en Colab):

Probad a jugar con el valor de SIZE_MULTIPLIER. Este parámetro simplemente modifica el tamaño de las imágenes de entrada, incrementando el número de operaciones necesarias y nos permite comprobar cómo se desenvuelve el hardware.

Rendimiento

Es necesario tener en consideración ciertos detalles para hacer un uso adecuado de las TPUs.

La ejecución sobre CPU y GPU no suele tener restricciones en cuanto al tamaño del batch, sin embargo en entornos de TPU este viene impuesto dada la arquitectura del hardware.

Una TPU como las que proporciona Google en su entorno Colab contiene 8 cores que funcionan como unidades independientes. Si queremos obtener el máximo de rendimiento debemos, por lo tanto, darle carga de trabajo a todos ellos.

Una buena alternativa es incrementar el batch size (en comparación a un entrenamiento con GPU), haciéndolo múltiplo del número de cores:

...
BATCH_SIZE = 64 * strategy.num_replicas_in_sync
...

Por otro lado, ser capaces de proporcionar datos a una velocidad adecuada es crítico para aprovechar al máximo su potencia. Existen varias recomendaciones:

• Si el dataset es pequeño, almacenarlo directamente en caché.
• Sino, almacenar los dataset en Buckets de Google Cloud Storage (GCS).
• Convertir los datos en formato TFRecord format y utilizar tf.data.TFRecordDataset para su lectura
• ...

Os dejo aquí el enlace con la guía de rendimiento oficial.

¿Vale para todo?

La respuesta corta sería que no. Si entramos en detalle y desarrollamos un poco esta respuesta, la conclusión sería: ¡por supuesto que no!

Como suele ser habitual, cada tecnología tiene su caso de uso y las TPUs no son una excepción.

Dada la arquitectura de las TPUs, su potencial destaca en tareas en las que el número de cálculos sobre matrices sea realmente elevado. Para cualquier otro tipo de tarea es muy posible que otras alternativas encajen mejor.

Google lo deja muy claro en la documentación, el rango de situaciones en las que se puede aprovechar esta tecnología es limitado.

A estas recomendaciones es necesario añadir aquellas situaciones en las que el esfuerzo de desarrollar pensando en una TPU no compense (proyectos de corta duración por ejemplo, o cargas de trabajo puntuales de tamaño moderado que no vayan a repetirse en el tiempo), ya que dadas las características del hardware, en ocasiones nos veremos obligados a alterar nuestro código (o incluso el dataset) para que encaje.

Conclusión

Las TPUs son una muestra de cómo la tecnología (tanto hardware como software) se adapta a nuevas necesidades y cómo esta nueva tecnología da pie a nuevos avances y oportunidades de forma que se genera un ciclo de retroalimentación mutua.

Los avances en ML están siendo imparables y es un área en el que precisamente los requisitos de cálculo son enormes y crecen día a día. Seguramente veamos mucha evolución en los próximos años en este tipo de hardware para dar respuesta a estas necesidades.