Otimizar a performance do TensorFlow usando o XProf

Este guia demonstra como usar as ferramentas disponíveis com o XProf para rastrear o desempenho dos modelos do TensorFlow no host (CPU), no dispositivo (GPU) ou em uma combinação de ambos.

A criação de perfis ajuda a entender o consumo de recursos de hardware (tempo e memória) das várias operações do TensorFlow (ops) no seu modelo e a resolver gargalos de desempenho e, por fim, fazer com que o modelo seja executado mais rápido.

Neste guia, mostramos como usar as várias ferramentas disponíveis e os diferentes modos de coleta de dados de performance do Profiler.

Se quiser criar um perfil da performance do modelo em Cloud TPUs, consulte o guia da Cloud TPU.

Coletar dados de performance

O XProf coleta atividades do host e rastreamentos de GPU do seu modelo do TensorFlow. É possível configurar o XProf para coletar dados de desempenho no modo programático ou de amostragem.

APIs de criação de perfil

Você pode usar as seguintes APIs para fazer o perfil:

• Modo programático usando o callback do TensorBoard Keras (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 usando a API tf.profiler Function

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

• Modo programático usando o gerenciador de contexto

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

• Modo de amostragem: faça o perfil sob demanda usando tf.profiler.experimental.server.start para iniciar um servidor gRPC com a execução do modelo do TensorFlow. Depois de iniciar o servidor gRPC e executar o modelo, você pode capturar um perfil usando o botão Capturar perfil no XProf. Use o script na seção "Instalar o criador de perfis" acima para iniciar uma instância do TensorBoard, caso ela ainda não esteja em execução.

  Por exemplo,

  # 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)
  

  Exemplo de criação de perfil para vários workers:

  # 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)
  

Use a caixa de diálogo Capturar perfil para especificar:

• Uma lista delimitada por vírgulas de URLs de serviço do perfil ou nomes de TPU.
• Uma duração de criação de perfil.
• O nível de rastreamento de chamadas de função do dispositivo, do host e do Python.
• Quantas vezes você quer que o Profiler tente capturar perfis se não tiver sucesso na primeira vez.

Como criar perfis de loops de treinamento personalizados

Para criar um perfil de loops de treinamento personalizados no seu código do TensorFlow, instrumente o loop de treinamento com a API tf.profiler.experimental.Trace para marcar os limites da etapa do XProf.

O argumento name é usado como prefixo para os nomes das etapas, o argumento de palavra-chave step_num é anexado aos nomes das etapas, e o argumento de palavra-chave _r faz com que esse evento de rastreamento seja processado como um evento de etapa pelo XProf.

Por exemplo,

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)

Isso vai ativar a análise de desempenho baseada em etapas do XProf e fazer com que os eventos de etapa apareçam no visualizador de rastreamento.

Inclua o iterador de conjunto de dados no contexto tf.profiler.experimental.Trace para uma análise precisa do pipeline de entrada.

O snippet de código abaixo é um antipadrão:

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 criação de perfil

O criador de perfis abrange vários casos de uso em quatro eixos diferentes. Algumas combinações já são aceitas, e outras serão adicionadas no futuro. Alguns dos casos de uso são:

• Criação de perfil local x remota: essas são duas maneiras comuns de configurar seu ambiente de criação de perfil. No criação de perfil local, a API de criação de perfil é chamada na mesma máquina em que o modelo está sendo executado, por exemplo, uma estação de trabalho local com GPUs. No criação de perfil remoto, a API de criação de perfil é chamada em uma máquina diferente de onde o modelo está sendo executado, por exemplo, em uma Cloud TPU.
• Criação de perfil de vários trabalhadores: é possível criar perfil de várias máquinas ao usar os recursos de treinamento distribuído do TensorFlow.
• Plataforma de hardware: crie perfis de CPUs, GPUs e TPUs.

A tabela abaixo oferece uma visão geral rápida dos casos de uso compatíveis com o TensorFlow mencionados acima:

Outros recursos

• O tutorial TensorFlow Profiler: criar perfil de desempenho do modelo com Keras e TensorBoard, em que você pode aplicar as dicas deste guia.
• A palestra Criação de perfil de desempenho no TensorFlow 2 da Conferência de desenvolvedores do TensorFlow 2020.
• A demonstração do TensorFlow Profiler da Conferência de desenvolvedores do TensorFlow 2020.