migration_debugging.ipynb

Kernel: Python 3

Copyright 2021 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Como depurar um pipeline de treinamento migrado do TensorFlow 2

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Este notebook demonstra como depurar um pipeline de treinamento ao migrar para o TensorFlow 2 (TF2). É composto pelos seguintes componentes:

1. Etapas sugeridas e amostras de código para depurar o pipeline de treinamento

2. Ferramentas para depuração

3. Outros recursos relacionados

Uma suposição é que você tem o código do TensorFlow 1 (TF1.x) e modelos treinados para comparação e deseja criar um modelo TF2 que alcance uma exatidão de validação semelhante.

Este notebook NÃO cobre problemas de desempenho de depuração para velocidade de inferência/treinamento, ou uso de memória.

Workflow de depuração

Um workflow geral para depurar seus pipelines de treinamento TF2 está mostrado abaixo. Observe que você não precisa seguir essas etapas nessa ordem. Você também pode usar uma abordagem de pesquisa binária onde testa o modelo numa etapa intermediária e depois restringe o escopo da depuração.

1. Corrigir erros de compilação e de tempo de execução

2. Validação única de passo para a frente (num guia separado)

   a. Num único dispositivo CPU

   • Verificar se as variáveis ​​são criadas apenas uma vez

   • Verificar a correspondência de contagens, nomes e formas de variáveis

   • Reiniciar todas as variáveis, verificar a equivalência numérica com toda a aleatoriedade desativada

   • Alinhar a geração de números aleatórios, verificar a equivalência numérica na inferência

   • (Opcional) Verificar se os checkpoints estão carregados corretamente e os modelos TF1.x/TF2 geram saída idêntica

   b. Num único dispositivo GPU/TPU

   c. Com estratégias para múltiplos dispositivos

3. Validação de equivalência numérica de treinamento de modelos para alguns passos (exemplos de código disponíveis abaixo)

   a. Validação de etapa de treinamento único usando dados pequenos e fixos num único dispositivo de CPU. Especificamente, verificar a equivalência numérica para os seguintes componentes

   • cálculo de perdas

   • métricas

   • taxa de aprendizagem

   • cálculo de gradiente e atualização

   b. Verificar as estatísticas depois do treinamento de 3 ou mais passos para verificar os comportamentos do otimizador, como o momento, ainda com dados fixos num único dispositivo de CPU

   c. Em um único dispositivo GPU/TPU

   d. Com estratégias para múltiplos dispositivos (veja a introdução de MultiProcessRunner no final deste artigo)

4. Teste de convergência de ponta a ponta em dataset real

   a. Verificar os comportamentos de treinamento com o TensorBoard

   • usar otimizadores simples, por exemplo, SGD e estratégias de distribuição simples, por exemplo, tf.distribute.OneDeviceStrategy primeiro

   • métricas de treinamento

   • métricas de avaliação

   • descobrir qual é a tolerância razoável para a aleatoriedade inerente

   b. Verificar a equivalência com otimizador avançado/agendador de taxa de aprendizado/estratégias de distribuição

   c. Verificar a equivalência ao usar precisão mista

5. Benchmarks de produto adicionais

Configuração

# The `DeterministicRandomTestTool` is only available from Tensorflow 2.8:
!pip install -q "tensorflow==2.9.*"

Validação única de passo para frente

A validação única de passo para frente, incluindo o carregamento do checkpoint, é abordada num colab diferente.

import sys
import unittest
import numpy as np

import tensorflow as tf
import tensorflow.compat.v1 as v1

Validação de equivalência numérica de treinamento de modelos para alguns passos

Defina a configuração do modelo e prepare um dataset falso.

params = {
    'input_size': 3,
    'num_classes': 3,
    'layer_1_size': 2,
    'layer_2_size': 2,
    'num_train_steps': 100,
    'init_lr': 1e-3,
    'end_lr': 0.0,
    'decay_steps': 1000,
    'lr_power': 1.0,
}

# make a small fixed dataset
fake_x = np.ones((2, params['input_size']), dtype=np.float32)
fake_y = np.zeros((2, params['num_classes']), dtype=np.int32)
fake_y[0][0] = 1
fake_y[1][1] = 1

step_num = 3

Defina o modelo TF1.x.

# Assume there is an existing TF1.x model using estimator API
# Wrap the model_fn to log necessary tensors for result comparison
class SimpleModelWrapper():
  def __init__(self):
    self.logged_ops = {}
    self.logs = {
        'step': [],
        'lr': [],
        'loss': [],
        'grads_and_vars': [],
        'layer_out': []}
     
  def model_fn(self, features, labels, mode, params):
      out_1 = tf.compat.v1.layers.dense(features, units=params['layer_1_size'])
      out_2 = tf.compat.v1.layers.dense(out_1, units=params['layer_2_size'])
      logits = tf.compat.v1.layers.dense(out_2, units=params['num_classes'])
      loss = tf.compat.v1.losses.softmax_cross_entropy(labels, logits)

      # skip EstimatorSpec details for prediction and evaluation 
      if mode == tf.estimator.ModeKeys.PREDICT:
          pass
      if mode == tf.estimator.ModeKeys.EVAL:
          pass
      assert mode == tf.estimator.ModeKeys.TRAIN

      global_step = tf.compat.v1.train.get_or_create_global_step()
      lr = tf.compat.v1.train.polynomial_decay(
        learning_rate=params['init_lr'],
        global_step=global_step,
        decay_steps=params['decay_steps'],
        end_learning_rate=params['end_lr'],
        power=params['lr_power'])
      
      optmizer = tf.compat.v1.train.GradientDescentOptimizer(lr)
      grads_and_vars = optmizer.compute_gradients(
          loss=loss,
          var_list=graph.get_collection(
              tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES))
      train_op = optmizer.apply_gradients(
          grads_and_vars,
          global_step=global_step)
      
      # log tensors
      self.logged_ops['step'] = global_step
      self.logged_ops['lr'] = lr
      self.logged_ops['loss'] = loss
      self.logged_ops['grads_and_vars'] = grads_and_vars
      self.logged_ops['layer_out'] = {
          'layer_1': out_1,
          'layer_2': out_2,
          'logits': logits}

      return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

  def update_logs(self, logs):
    for key in logs.keys():
      model_tf1.logs[key].append(logs[key])

A seguinte classe v1.keras.utils.DeterministicRandomTestTool fornece um gerenciador de contexto scope() que pode fazer com que operações aleatórias stateful usem a mesma semente em ambos os grafos/sessões TF1 e execução antecipada (eager).

A ferramenta fornece dois modos de teste:

1. constant que usa a mesma semente para cada operação, não importa quantas vezes tenha sido chamada e,

2. num_random_ops que usa o número de operações stateful aleatórias observadas anteriormente como a semente da operação.

Isto se aplica tanto às operações aleatórias stateful usadas para criar e inicializar variáveis ​​quanto às operações aleatórias stateful usadas no cálculo (como para camadas de dropout).

random_tool = v1.keras.utils.DeterministicRandomTestTool(mode='num_random_ops')

Execute o modelo TF1.x no modo grafo. Colete estatísticas para as 3 primeiras etapas de treinamento para comparação de equivalência numérica.

with random_tool.scope():
  graph = tf.Graph()
  with graph.as_default(), tf.compat.v1.Session(graph=graph) as sess:
    model_tf1 = SimpleModelWrapper()
    # build the model
    inputs = tf.compat.v1.placeholder(tf.float32, shape=(None, params['input_size']))
    labels = tf.compat.v1.placeholder(tf.float32, shape=(None, params['num_classes']))
    spec = model_tf1.model_fn(inputs, labels, tf.estimator.ModeKeys.TRAIN, params)
    train_op = spec.train_op

    sess.run(tf.compat.v1.global_variables_initializer())
    for step in range(step_num):
      # log everything and update the model for one step
      logs, _ = sess.run(
          [model_tf1.logged_ops, train_op],
          feed_dict={inputs: fake_x, labels: fake_y})
      model_tf1.update_logs(logs)

Defina o modelo TF2.

class SimpleModel(tf.keras.Model):
  def __init__(self, params, *args, **kwargs):
    super(SimpleModel, self).__init__(*args, **kwargs)
    # define the model
    self.dense_1 = tf.keras.layers.Dense(params['layer_1_size'])
    self.dense_2 = tf.keras.layers.Dense(params['layer_2_size'])
    self.out = tf.keras.layers.Dense(params['num_classes'])
    learning_rate_fn = tf.keras.optimizers.schedules.PolynomialDecay(
      initial_learning_rate=params['init_lr'],
      decay_steps=params['decay_steps'],
      end_learning_rate=params['end_lr'],
      power=params['lr_power'])  
    self.optimizer = tf.keras.optimizers.legacy.SGD(learning_rate_fn)
    self.compiled_loss = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
    self.logs = {
        'lr': [],
        'loss': [],
        'grads': [],
        'weights': [],
        'layer_out': []}

  def call(self, inputs):
    out_1 = self.dense_1(inputs)
    out_2 = self.dense_2(out_1)
    logits = self.out(out_2)
    # log output features for every layer for comparison
    layer_wise_out = {
        'layer_1': out_1,
        'layer_2': out_2,
        'logits': logits}
    self.logs['layer_out'].append(layer_wise_out)
    return logits

  def train_step(self, data):
    x, y = data
    with tf.GradientTape() as tape:
      logits = self(x)
      loss = self.compiled_loss(y, logits)
    grads = tape.gradient(loss, self.trainable_weights)
    # log training statistics
    step = self.optimizer.iterations.numpy()
    self.logs['lr'].append(self.optimizer.learning_rate(step).numpy())
    self.logs['loss'].append(loss.numpy())
    self.logs['grads'].append(grads)
    self.logs['weights'].append(self.trainable_weights)
    # update model
    self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
    return

Execute o modelo TF2 no modo eager. Colete estatísticas para as 3 primeiras etapas de treinamento para comparação de equivalência numérica.

random_tool = v1.keras.utils.DeterministicRandomTestTool(mode='num_random_ops')
with random_tool.scope():
  model_tf2 = SimpleModel(params)
  for step in range(step_num):
    model_tf2.train_step([fake_x, fake_y])

Compare a equivalência numérica para os primeiros passos de treinamento.

Você também pode dar uma olhada no notebook Validando exatidão e equivalência numérica para conselhos adicionais sobre equivalência numérica.

np.testing.assert_allclose(model_tf1.logs['lr'], model_tf2.logs['lr'])
np.testing.assert_allclose(model_tf1.logs['loss'], model_tf2.logs['loss'])
for step in range(step_num):
  for name in model_tf1.logs['layer_out'][step]:
    np.testing.assert_allclose(
        model_tf1.logs['layer_out'][step][name],
        model_tf2.logs['layer_out'][step][name])

Testes de unidade

Existem alguns tipos de teste de unidade que podem ajudar a depurar seu código de migração.

1. Validação única de passo para frente

2. Validação de equivalência numérica de treinamento de modelos para alguns passos

3. Benchmark de desempenho de inferência

4. O modelo treinado faz previsões corretas em pontos de dados fixos e simples

Você pode usar @parameterized.parameters para testar modelos com diferentes configurações. Detalhes com amostra de código.

Veja que é possível executar APIs de sessão e execução antecipada (eager) no mesmo caso de teste. Os trechos de código abaixo mostram como.

import unittest

class TestNumericalEquivalence(unittest.TestCase):

  # copied from code samples above
  def setup(self):
    # record statistics for 100 training steps
    step_num = 100

    # setup TF 1 model
    random_tool = v1.keras.utils.DeterministicRandomTestTool(mode='num_random_ops')
    with random_tool.scope():
      # run TF1.x code in graph mode with context management
      graph = tf.Graph()
      with graph.as_default(), tf.compat.v1.Session(graph=graph) as sess:
        self.model_tf1 = SimpleModelWrapper()
        # build the model
        inputs = tf.compat.v1.placeholder(tf.float32, shape=(None, params['input_size']))
        labels = tf.compat.v1.placeholder(tf.float32, shape=(None, params['num_classes']))
        spec = self.model_tf1.model_fn(inputs, labels, tf.estimator.ModeKeys.TRAIN, params)
        train_op = spec.train_op

        sess.run(tf.compat.v1.global_variables_initializer())
        for step in range(step_num):
          # log everything and update the model for one step
          logs, _ = sess.run(
              [self.model_tf1.logged_ops, train_op],
              feed_dict={inputs: fake_x, labels: fake_y})
          self.model_tf1.update_logs(logs)

    # setup TF2 model
    random_tool = v1.keras.utils.DeterministicRandomTestTool(mode='num_random_ops')
    with random_tool.scope():
      self.model_tf2 = SimpleModel(params)
      for step in range(step_num):
        self.model_tf2.train_step([fake_x, fake_y])
  
  def test_learning_rate(self):
    np.testing.assert_allclose(
        self.model_tf1.logs['lr'],
        self.model_tf2.logs['lr'])

  def test_training_loss(self):
    # adopt different tolerance strategies before and after 10 steps
    first_n_step = 10

    # absolute difference is limited below 1e-5
    # set `equal_nan` to be False to detect potential NaN loss issues
    abosolute_tolerance = 1e-5
    np.testing.assert_allclose(
        actual=self.model_tf1.logs['loss'][:first_n_step],
        desired=self.model_tf2.logs['loss'][:first_n_step],
        atol=abosolute_tolerance,
        equal_nan=False)
    
    # relative difference is limited below 5%
    relative_tolerance = 0.05
    np.testing.assert_allclose(self.model_tf1.logs['loss'][first_n_step:],
                               self.model_tf2.logs['loss'][first_n_step:],
                               rtol=relative_tolerance,
                               equal_nan=False)

Ferramentas de depuração

tf.print

tf.print vs print/logging.info

• Com argumentos configuráveis, tf.print pode exibir recursivamente os primeiros e últimos elementos de cada dimensão para tensores impressos. Veja a Documentação da API para mais detalhes.

• Para execução antecipada (eager), print e tf.print imprimem o valor do tensor. Mas print pode envolver a cópia do dispositivo para o host, o que poderá desacelerar seu código.

• Para o modo grafo, incluindo o uso dentro tf.function, você precisa usar tf.print para imprimir o valor real do tensor. tf.print é compilado em um op no grafo, enquanto print e logging.info registram apenas no tempo de rastreamento, o que geralmente não é o que você deseja.

• tf.print também suporta a impressão de tensores compostos como tf.RaggedTensor e tf.sparse.SparseTensor.

• Você também pode usar um callback para monitorar métricas e variáveis. Verifique como usar callbacks personalizados com logs dict e com o atributo self.model.

tf.print vs print dentro de tf.function

# `print` prints info of tensor object
# `tf.print` prints the tensor value
@tf.function
def dummy_func(num):
  num += 1
  print(num)
  tf.print(num)
  return num

_ = dummy_func(tf.constant([1.0]))

# Output:
# Tensor("add:0", shape=(1,), dtype=float32)
# [2]

tf.distribute.Strategy

• Se a tf.function contendo tf.print for executada nos workers, por exemplo, ao usar TPUStrategy ou ParameterServerStrategy, será necessário verificar os logs do worker/servidor de parâmetros para encontrar os valores impressos.

• Para print ou logging.info , os logs serão impressos no coordenador ao usar ParameterServerStrategy e os logs serão impressos no STDOUT em worker0 ao usar TPUs.

tf.keras.Model

• Ao usar modelo das APIs Sequential e Functional, se você deseja imprimir valores, por exemplo, entradas de modelo ou características intermediárias após algumas camadas, você pode dispor das seguintes alternativas.

  1. Escrever uma camada personalizada que use tf.print para imprimir as entradas.

  2. Incluir as saídas intermediárias que deseja inspecionar nas saídas do modelo.

• As camadas tf.keras.layers.Lambda têm limitações de (des)serialização. Para evitar problemas de carregamento de checkpoint, escreva uma camada de personalizada numa subclasse. Veja a documentação da API para mais detalhes.

• Você não pode usar tf.print com saídas intermediárias em um tf.keras.callbacks.LambdaCallback se não tiver acesso aos valores reais, mas apenas aos objetos simbólicos do tensor Keras.

Opção 1: escrever uma camada personalizada

class PrintLayer(tf.keras.layers.Layer):
  def call(self, inputs):
    tf.print(inputs)
    return inputs

def get_model():
  inputs = tf.keras.layers.Input(shape=(1,))
  out_1 = tf.keras.layers.Dense(4)(inputs)
  out_2 = tf.keras.layers.Dense(1)(out_1)
  # use custom layer to tf.print intermediate features
  out_3 = PrintLayer()(out_2)
  model = tf.keras.Model(inputs=inputs, outputs=out_3)
  return model

model = get_model()
model.compile(optimizer="adam", loss="mse")
model.fit([1, 2, 3], [0.0, 0.0, 1.0])

Opção 2: incluir as saídas intermediárias que deseja inspecionar nas saídas do modelo.

Observe que, nesse caso, você poderá precisar de algumas personalizações para usar Model.fit.

def get_model():
  inputs = tf.keras.layers.Input(shape=(1,))
  out_1 = tf.keras.layers.Dense(4)(inputs)
  out_2 = tf.keras.layers.Dense(1)(out_1)
  # include intermediate values in model outputs
  model = tf.keras.Model(
      inputs=inputs,
      outputs={
          'inputs': inputs,
          'out_1': out_1,
          'out_2': out_2})
  return model

pdb

Você pode usar o pdb tanto no terminal como no Colab para inspecionar valores intermediários para depuração.

Visualização do grafo com o TensorBoard

Você pode examinar o grafo do TensorFlow com o TensorBoard. O TensorBoard também é suportado no colab. O TensorBoard é uma ótima ferramenta para visualizar resumos. Você pode usá-lo para comparar a taxa de aprendizado, pesos do modelo, escala de gradiente, métricas de treinamento/validação ou até mesmo modelar resultados intermediários entre o modelo TF1.x e o modelo TF2 migrado por meio do processo de treinamento e ver se os valores são os esperados.

TensorFlow Profiler

O TensorFlow Profiler pode ajudar você a visualizar a linha do tempo de execução em GPUs/TPUs. Veja o uso básico nesta demonstração do Colab.

MultiProcessRunner

O MultiProcessRunner é uma ferramenta útil ao depurar com MultiWorkerMirroredStrategy e ParameterServerStrategy. Você pode dar uma olhada neste exemplo concreto para ver como usar.

Especificamente, para essas duas estratégias, é recomendável 1) não depender apenas de testes de unidade para cobrir seu fluxo, 2) mas também tentar reproduzir falhas usando-o num teste de unidade para evitar que uma tarefa distribuída real seja executada toda vez que houver uma tentativa de consertar a falha.