Optimiza el rendimiento de TensorFlow con XProf

En esta guía, se muestra cómo usar las herramientas disponibles con XProf para hacer un seguimiento del rendimiento de tus modelos de TensorFlow en el host (CPU), el dispositivo (GPU) o en una combinación de ambos.

La generación de perfiles ayuda a comprender el consumo de recursos de hardware (tiempo y memoria) de las diversas operaciones (ops) de TensorFlow en tu modelo, resolver los cuellos de botella de rendimiento y, en última instancia, hacer que el modelo se ejecute más rápido.

En esta guía, se explica cómo usar las diversas herramientas disponibles y los diferentes modos en que Profiler recopila datos de rendimiento.

Si deseas generar un perfil del rendimiento de tu modelo en Cloud TPU, consulta la guía de Cloud TPU.

Recopila datos de rendimiento

XProf recopila las actividades del host y los registros de GPU de tu modelo de TensorFlow. Puedes configurar XProf para recopilar datos de rendimiento a través del modo programático o del modo de muestreo.

APIs de generación de perfiles

Puedes usar las siguientes APIs para realizar la generación de perfiles.

• Modo programático con la devolución de llamada de Keras de TensorBoard (tf.keras.callbacks.TensorBoard)

  # Profile from batches 10 to 15
  tb_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir,
                                               profile_batch='10, 15')
  
  # Train the model and use the TensorBoard Keras callback to collect
  # performance profiling data
  model.fit(train_data,
            steps_per_epoch=20,
            epochs=5,
            callbacks=[tb_callback])
  

• Modo programático con la API de Functions tf.profiler

  tf.profiler.experimental.start('logdir')
  # Train the model here
  tf.profiler.experimental.stop()
  

• Modo programático con el administrador de contexto

  with tf.profiler.experimental.Profile('logdir'):
      # Train the model here
      pass
  

• Modo de muestreo: Realiza la creación de perfiles a pedido con tf.profiler.experimental.server.start para iniciar un servidor de gRPC con la ejecución de tu modelo de TensorFlow. Después de iniciar el servidor de gRPC y ejecutar tu modelo, puedes capturar un perfil con el botón Capture Profile en XProf. Usa la secuencia de comandos de la sección Install profiler que se encuentra más arriba para iniciar una instancia de TensorBoard si aún no se está ejecutando.

  Por ejemplo:

  # Start a profiler server before your model runs.
  tf.profiler.experimental.server.start(6009)
  # (Model code goes here).
  #  Send a request to the profiler server to collect a trace of your model.
  tf.profiler.experimental.client.trace('grpc://localhost:6009',
                                        'gs://your_tb_logdir', 2000)
  

  Ejemplo para generar perfiles de varios trabajadores:

  # E.g., your worker IP addresses are 10.0.0.2, 10.0.0.3, 10.0.0.4, and you
  # would like to profile for a duration of 2 seconds.
  tf.profiler.experimental.client.trace(
      'grpc://10.0.0.2:8466,grpc://10.0.0.3:8466,grpc://10.0.0.4:8466',
      'gs://your_tb_logdir',
      2000)
  

Usa el cuadro de diálogo Capture Profile para especificar lo siguiente:

• Es una lista delimitada por comas de URLs del servicio de perfil o nombres de TPU.
• Es la duración del registro de perfil.
• Es el nivel de registro de llamadas a funciones de Python, host y dispositivo.
• Cantidad de veces que deseas que Profiler vuelva a intentar capturar perfiles si no lo logra la primera vez.

Cómo crear perfiles de bucles de entrenamiento personalizados

Para generar perfiles de bucles de entrenamiento personalizados en tu código de TensorFlow, instrumenta el bucle de entrenamiento con la API de tf.profiler.experimental.Trace para marcar los límites de los pasos de XProf.

El argumento name se usa como prefijo para los nombres de los pasos, el argumento de palabra clave step_num se agrega a los nombres de los pasos y el argumento de palabra clave _r hace que XProf procese este evento de registro como un evento de paso.

Por ejemplo:

for step in range(NUM_STEPS):
    with tf.profiler.experimental.Trace('train', step_num=step, _r=1):
        train_data = next(dataset)
        train_step(train_data)

Esto habilitará el análisis de rendimiento basado en pasos de XProf y hará que los eventos de pasos aparezcan en el lector de seguimiento.

Asegúrate de incluir el iterador del conjunto de datos en el contexto de tf.profiler.experimental.Trace para realizar un análisis preciso de la canalización de entrada.

El siguiente fragmento de código es un antipatrón:

for step, train_data in enumerate(dataset):
    with tf.profiler.experimental.Trace('train', step_num=step, _r=1):
        train_step(train_data)

Casos de uso de generación de perfiles

El generador de perfiles abarca varios casos de uso en cuatro ejes diferentes. Actualmente, se admiten algunas combinaciones y, en el futuro, se agregarán otras. Estos son algunos de los casos de uso:

• Perfilamiento local vs. remoto: Estas son dos formas comunes de configurar tu entorno de perfilamiento. En el perfilamiento local, se llama a la API de perfilamiento en la misma máquina en la que se ejecuta tu modelo, por ejemplo, una estación de trabajo local con GPUs. En la generación de perfiles remota, se llama a la API de generación de perfiles en una máquina diferente de la que ejecuta tu modelo, por ejemplo, en una Cloud TPU.
• Generación de perfiles de varios trabajadores: Puedes generar perfiles de varias máquinas cuando usas las capacidades de entrenamiento distribuido de TensorFlow.
• Plataforma de hardware: Genera perfiles de CPUs, GPUs y TPUs.

En la siguiente tabla, se proporciona una descripción general rápida de los casos de uso compatibles con TensorFlow que se mencionaron anteriormente:

Recursos adicionales

• El instructivo TensorFlow Profiler: Cómo generar perfiles del rendimiento del modelo con Keras y TensorBoard, en el que puedes aplicar los consejos de esta guía.
• La charla Performance profiling in TensorFlow 2 de la Cumbre de desarrolladores de TensorFlow 2020.
• La demostración de TensorFlow Profiler de la Cumbre de desarrolladores de TensorFlow 2020.