Guía de rendimiento

Cloud TPU proporciona alto rendimiento a bajo costo. A fin de aumentar aún más el rendimiento de Cloud TPU, puedes ajustar los parámetros de configuración para tu aplicación, como también identificar y resolver cualquier cuello de botella que limite el rendimiento.

Esta guía te ayuda a maximizar el rendimiento de Cloud TPU; para ello, te enseña a hacer lo siguiente:

Además, los recursos que aparecen a continuación describen cómo hacer lo siguiente:

Rendimiento del procesamiento de modelos

En esta sección, se describen problemas generales que pueden reducir el rendimiento del modelo y cómo puedes abordarlos.

  1. Dedicar demasiado tiempo al procesamiento previo de datos

    La pila de software de TensorFlow-TPU permite que el usuario realice un procesamiento previo de datos complejo en la CPU antes de ingresar datos en la TPU. Dado que la TPU es tan rápida, el procesamiento de datos de entrada puede convertirse fácilmente en un cuello de botella si los datos de entrada son muy grandes o se procesan de manera compleja.

    Las herramientas de generación de perfiles de Cloud TPU proporcionan una manera simple de medir si el procesamiento de entrada es un cuello de botella en tu modelo. Si este es el caso, realiza un procesamiento previo sin conexión como un costo único, en lugar de incurrir en el costo de cada ciclo de entrenamiento.

  2. Tamaño del lote demasiado pequeño debido a la fragmentación (división de lotes en los núcleos)

    Si usas un tamaño de lote de entrenamiento de n, se divide o fragmenta en 8 núcleos de TPU. Si n = 128 y comparas una TPU con una P100 única, la GPU podría ejecutarse al 70% de los FLOPS máximos, mientras que la TPU solo se ejecuta al 20% de los FLOPS máximos (porque el fragmentación es El tamaño del lote de TPU es, en realidad, 16, 128 (28 / 8) y no en 128. En este ejemplo, un tamaño de lote de TPU de 1,024 usaría mejor la memoria de TPU.

    Para optimizar el uso de la memoria, usa el tamaño del lote más grande que cabe en la memoria. Cada núcleo de TPU usa una matriz de celdas de memoria de 128 x 128. En general, tu tamaño de lote debe ser divisible por 128 para usar la memoria de TPU de manera más efectiva.

Rendimiento del compilador XLA

XLA es un compilador para el aprendizaje automático que puede producir binarios para TPU, CPU, GPU y otras plataformas. Es parte de la base de código estándar de TensorFlow. Los modelos de TensorFlow para Cloud TPU se convierten en un grafo XLA, que, a su vez, XLA compila en una TPU ejecutable. Si deseas conocer más detalles sobre cómo XLA y TensorFlow interactúan, consulta Descripción general de XLA.

El hardware de Cloud TPU difiere de las CPU y GPU. En un nivel alto, las CPU se pueden caracterizar por tener una cantidad baja de subprocesos de rendimiento alto. Las GPU se pueden caracterizar por tener una gran cantidad de subprocesos de rendimiento bajo. Cloud TPU, con su unidad de matriz de 128 x 128, puede considerarse un único subproceso muy poderoso, que puede realizar 16,000 ops por ciclo, o como subprocesos pequeños de 128 x 128 conectados por la canalización. De la misma manera, cuando se trata la memoria, se buscan múltiplos de 8 (flotantes), como también múltiplos de 128 para operaciones orientadas a la unidad de la matriz.

Consecuencias de los mosaicos

Los arreglos en Cloud TPU están en mosaicos. Esto implica el relleno de una de las dimensiones a un múltiplo de 8 y de otra a un múltiplo de 128. XLA realiza transformaciones del diseño de los datos, y estos se organizan en la memoria de modo que el hardware pueda procesarlos de forma eficiente. La heurística impulsa estas transformaciones. Si bien son beneficiosas en la mayoría de los casos, siempre existe la posibilidad de que el compilador haga algo incorrecto. Para lograr el máximo rendimiento, puede ser beneficioso experimentar con diferentes configuraciones de modelos.

Regla general: minimiza el costo del relleno

Una consecuencia de utilizar el esquema de memoria en mosaico es que la utilización eficiente de la memoria depende de la cantidad de memoria gastada en la sobrecarga de relleno. Para utilizar Cloud TPU de la forma más eficiente, usa tamaños de dimensión que minimicen la sobrecarga de los mosaicos.

Por ejemplo, un resultado de convolución tiene (1) lote, (2) atributos de resultados y (3) dimensiones espaciales de resultados. Una de las dimensiones, ya sea de lote o del atributo de resultado, se rellenará a un múltiplo de 8, mientras que la otra lo hará a uno de 128. Las dimensiones espaciales de resultados no se rellenarán.

En general, para una operación con dimensiones espaciales (tf.nn.pool, tf.conv2d, etc.), las dimensiones espaciales nunca se rellenan.

Regla general: elige valores eficientes para dimensiones de lotes y atributos

Las dimensiones de lotes y de atributos están sujetas al relleno; ten cuidado cuando determines los tamaños del lote y del atributo. Para aprovechar la unidad de matriz de 128 x 128, busca valores razonablemente grandes (>=128) de lotes o atributos, preferentemente ambos.

En la mayoría de los casos, un tamaño de lote de 128 es suficiente para que la unidad de matriz de Cloud TPU se mantenga ocupada. La ejecución de tamaños de lote mayores también funciona en Cloud TPU, pero se recomienda utilizar un tamaño de lote que sea múltiplo de 128. Si tu modelo no puede funcionar con esa configuración, intenta utilizar un tamaño de lote que sea múltiplo de 8 a fin de minimizar el impacto del relleno.

De modo similar, las dimensiones de atributos también se asignan a una dimensión que se rellena, ya sea con 8 o 128, en función de las decisiones tomadas por el algoritmo de diseño de XLA. Esto significa que la dimensión del atributo desperdicia la menor cantidad de espacio en múltiplos de 8 o 128.

Fusión

La fusión es una técnica general que utiliza el compilador de XLA a fin de optimizar programas. Una operación fusionada es la combinación de múltiples operaciones constituyentes que se ejecutarán de forma conjunta.

Por ejemplo, considera las siguientes series de operaciones:

tmp = tf.add(x, y)
result = tf.multiply(tmp, z)

Este código es, aproximadamente, equivalente al siguiente pseudocódigo:

for (i = 0; i < element_count; i++) {
  tmp[i] = x[i] + y[i];
}

for (i = 0; i < element_count; i++) {
  result = tmp[i] * z[i];
}

Con la fusión, los accesos al arreglo suceden al mismo tiempo:

for (i = 0; i < element_count; i++) {
  result = (x[i] + y[i]) * z[i];
}

En este ejemplo, la cantidad de viajes ida y vuelta a la memoria se reduce, y XLA no necesita asignar más espacio para 'tmp'.

La fusión es una optimización crítica y beneficia a Cloud TPU de diferentes maneras:

  • Reduce las transferencias de memoria, ya que quita la necesidad de almacenar resultados inmediatos en la memoria principal, lo cual es lento.
  • Permite una mejor utilización de unidades de hardware, que, de otra manera, no se hubieran utilizado.
  • Puede reducir la utilización de memoria de un modelo, ya que se necesitan menos búferes al mismo tiempo.

Transmisión

La transmisión se produce implícitamente cuando se combinan dos tensores con formas diferentes pero compatibles.

Por ejemplo, tf.add(vector, matrix) requiere que el vector se transmita a la forma de la matriz. El resultado de la operación tiene la misma forma que la matriz. Si deseas obtener más detalles, consulta la guía transmisión de arreglos.

Si bien las transmisiones, normalmente, pueden fusionarse con sus consumidores, cuando se obliga a que una transmisión se materialice, se obtiene como resultado un rendimiento bajo y una mayor presión en la memoria.

En el siguiente ejemplo, la transmisión implícita en la adición de un vector y una matriz no se puede fusionar con el argmax, lo que da como resultado una transmisión materializada:

tf.argmax(tf.add(vector, zero_matrix), axis=0)

Recomendaciones específicas de la función de TensorFlow

Consulta la lista completa de operaciones de TensorFlow disponibles en Cloud TPU.

tf.matmul

  • La transposición del resultado de cualquiera de los operandos es efectivamente gratuita.
  • Ten en cuenta que tf.matmul admite la fusión en su entrada y salida. Esto significa que las funciones de activación o los sesgos aplicados directamente al resultado de tf.matmul tienen una sobrecarga baja.

tf.nn.conv_n_d, tf.nn.depthwise_conv2d, tf.nn.separable_conv2d

  • Para las activaciones, las dimensiones de lotes y de atributos se rellenan a un múltiplo de 8 o 128.
  • Primero, XLA hace un seguimiento del tamaño más común de las dimensiones de los lotes para las convoluciones en el módulo. Esto ayuda a distinguir entre convoluciones hacia delante, de gradiente de activación y de gradiente de kernel.
  • Si el tamaño más común de lote es mayor o igual que 64, se da lo siguiente:
    • El lote se rellena a un múltiplo de 128 y el atributo a uno de 8 para las convoluciones hacia delante y atrás.
    • El lote se rellena a un múltiplo de 8 y el atributo a uno de 128 para las convoluciones de actualización de gradiente.
  • Si el tamaño más común de lote es menor que 64, se da lo siguiente:
    • El lote se rellena a un múltiplo de 8 y el atributo a uno de 128 para las convoluciones hacia delante y atrás.
    • El lote se rellena a un múltiplo de 128 y el atributo a uno de 8 para las convoluciones de actualización de gradiente.
    • La transposición de las activaciones justo antes de enviarlas a la convolución es gratis si solo intercambia las dimensiones del lote y del atributo de entrada.
  • Para el kernel, las dimensiones del atributo de entrada y del resultado se rellenan a un múltiplo de 8 o 128. La determinación exacta está influenciada por los productores y otros consumidores del kernel.
    • La transposición del kernel justo antes de enviarla a la convolución es gratis si solo intercambia las dimensiones del atributo de entrada y del resultado.
  • Para el resultado, las dimensiones del lote y el atributo se rellenan a un múltiplo de 8 o 128.
  • La transposición del resultado de una convolución es gratis si solo intercambia las dimensiones del atributo de entrada y del lote.
  • Ten en cuenta que tf.nn.conv_n_d admite la fusión en su resultado, las activaciones o el kernel. Esto significa que las funciones de activación o los sesgos aplicados directamente en el resultado tienen una sobrecarga baja.
  • tf.nn.avg_pool, tf.nn.max_pool

    • Se aplican las reglas de relleno: las dimensiones espaciales son más importantes que el lote y el atributo. Cada lote y atributo se puede rellenar a un múltiplo de 8 o 128.
    • Por lo general, el diseño de una operación de grupo coincide con las convoluciones que entran o salen de esta.
    • El cálculo del gradiente para tf.nn.max_pool puede ser más lento que el equivalente a tf.nn.avg_pool. Considera cambiar de la agrupación máxima a la promedio cuando sea posible.

    tf.concat, tf.slice, tf.strided_slice

    • Evita los segmentos y las concatenaciones innecesarias. En una dimensión que se ha rellenado, estas se consideran más costosas.
      • El movimiento de datos se minimiza si la dimensión de los segmentos no tiene sobrecarga de relleno.

    tf.transpose

    • La transposición de cualquiera de los operandos de tf.matmul o su resultado es gratuita.
    • La transposición de las activaciones de un tf.conv_n_d es gratuita si se intercambian las dimensiones del lote y del atributo de resultado.
    • La transposición del kernel de un tf.conv_n_d es gratuita si se intercambian las dimensiones de los atributos de entrada y salida.
    • La transposición del resultado de tf.conv_n_d es gratuita si intercambia las dimensiones de lote y del atributo del resultado.

    tf.batch_to_space, tf.space_to_batch, tf.space_to_depth, tf.depth_to_space

    • Estas son costosas, ya que implican mover datos de dimensiones rellenadas a otras no rellenadas, y viceversa.

    tf.reshape

    • El cambio de forma podría ser costoso en Cloud TPU cuando se mueven en torno a datos en una dimensión rellenada.
    • Puede ser beneficioso cambiar la forma de los datos a R1 en el host y volver a cambiarla a una forma de dimensión mayor en el dispositivo si existe un relleno sustancial. Esto puede lograr que las transferencias entre el host y el dispositivo sean más eficientes.
      • Además, puede ayudar con la utilización máxima de la memoria, ya que el parámetro empaquetado puede desempaquetarse bajo demanda.

    tf.random_uniform, tf.distributions.Bernoulli, tf.random_normal, tf.multinomial

    • La generación de números pseudo al azar para distribuciones uniformes o de Bernoulli es muy rápida.
    • Las distribuciones normales son un poco más costosas que las uniformes o de Bernoulli.
    • La generación de números pseudo al azar para distribuciones categóricas o multinominales es considerablemente más costosa.

    tf.reduce_all, tf.reduce_any, tf.reduce_logsumexp, tf.reduce_max, tf.reduce_min, tf.reduce_prod, tf.reduce_sum

    • Se pueden realizar reducciones múltiples con la misma forma de entrada y resultado en paralelo mediante la fusión.
    • Vuelve a escribir cadenas secuenciales de reducciones en otras paralelas, de ser posible.
  • Las reducciones admiten la fusión de operaciones relativas a los elementos en su entrada, pero no en sus resultados. Cuando sea posible, vuelve a escribir expresiones para promover la fusión. Por ejemplo:

        tf.multiply(tf.reduce_sum(x), y)
    

    En:

        tf.reduce_sum(tf.multiply(x, y))
    
  • tf.nn.batch_normalization, tf.nn.fused_batch_norm, tf.layers.batch_normalization

    • El compilador de Cloud TPU puede reducir de forma eficiente las variantes fusionadas de TensorFlow de la normalización de lotes. Su uso puede ser considerablemente más eficiente que la alternativa.
      • Es preferible tf.nn.fused_batch_norm en lugar de tf.nn.batch_normalization.
      • En tf.layers.batch_normalization, configura el argumento "fusionado" como verdadero.

    tf.nn.depthwise_conv2d, tf.nn.separable_conv2d

    • Estos aún no están del todo optimizados y podrían tener un rendimiento peor que una convolución normal equivalente.