Depurar pipeline de treinamento migrado do TF2

Veja no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

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

  1. Etapas sugeridas e exemplos de código para depurar o pipeline de treinamento
  2. Ferramentas para depuração
  3. Outros recursos relacionados

Uma suposição é que você tem código TF1.xe modelos treinados para comparação e deseja construir um modelo TF2 que alcance precisão de validação semelhante.

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

Fluxo de trabalho de depuração

Abaixo está um fluxo de trabalho geral para depurar seus pipelines de treinamento do TF2. Observe que você não precisa seguir essas etapas em ordem. Você também pode usar uma abordagem de pesquisa binária em que testa o modelo em uma etapa intermediária e restringe o escopo de depuração.

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

  2. Validação de passagem única (em um guia separado)

    uma. Em um único dispositivo de CPU

    • Verifique se as variáveis ​​são criadas apenas uma vez
    • Verifique a correspondência de contagens de variáveis, nomes e formas
    • Redefina todas as variáveis, verifique a equivalência numérica com toda aleatoriedade desabilitada
    • Alinhar a geração de números aleatórios, verificar a equivalência numérica na inferência
    • (Opcional) Checkpoints são carregados corretamente e modelos TF1.x/TF2 geram saída idêntica

    b. Em um único dispositivo GPU/TPU

    c. Com estratégias para vários dispositivos

  3. Validação de equivalência numérica de treinamento de modelo para algumas etapas (amostras de código disponíveis abaixo)

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

    • cálculo de perdas
    • Métricas
    • taxa de Aprendizagem
    • cálculo e atualização de gradiente

    b. Verifique as estatísticas após treinar 3 ou mais etapas para verificar os comportamentos do otimizador, como o momento, ainda com dados fixos em um único dispositivo de CPU

    c. Em um único dispositivo GPU/TPU

    d. Com estratégias para vários dispositivos (verifique a introdução do MultiProcessRunner na parte inferior)

  4. Testes de cobertura de ponta a ponta em um conjunto de dados real

    uma. Verifique os comportamentos de treinamento com o TensorBoard

    • use 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. Verifique a equivalência com o otimizador avançado/programador de taxa de aprendizado/estratégias de distribuição

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

  5. Referências adicionais de produtos

Configurar

pip uninstall -y -q tensorflow

# Install tf-nightly as the DeterministicRandomTestTool is only available in
# Tensorflow 2.8
pip install -q tf-nightly

Validação de passagem única

A validação de passagem única, incluindo o carregamento do ponto de verificação, é abordada em uma 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 modelo para algumas etapas

Configure a configuração do modelo e prepare um conjunto de dados 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 classe v1.keras.utils.DeterministicRandomTestTool a seguir fornece um scope() que pode fazer com que as operações aleatórias com estado usem a mesma semente em ambos os gráficos/sessões TF1 e execução antecipada,

A ferramenta oferece dois modos de teste:

  1. constant que usa a mesma semente para cada operação, não importa quantas vezes ela tenha sido chamada e,
  2. num_random_ops que usa o número de operações aleatórias com estado observadas anteriormente como a semente da operação.

Isso se aplica tanto às operações aleatórias com estado usadas para criar e inicializar variáveis ​​quanto às operações aleatórias com estado usadas na computação (como para camadas de descarte).

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

WARNING:tensorflow:From /tmp/ipykernel_26769/2689227634.py:1: The name tf.keras.utils.DeterministicRandomTestTool is deprecated. Please use tf.compat.v1.keras.utils.DeterministicRandomTestTool instead.

Execute o modelo TF1.x no modo gráfico. Colete estatísticas para as primeiras 3 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)

/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel_launcher.py:14: UserWarning: `tf.layers.dense` is deprecated and will be removed in a future version. Please use `tf.keras.layers.Dense` instead.
  
/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/legacy_tf_layers/core.py:261: UserWarning: `layer.apply` is deprecated and will be removed in a future version. Please use `layer.__call__` method instead.
  return layer.apply(inputs)
/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel_launcher.py:15: UserWarning: `tf.layers.dense` is deprecated and will be removed in a future version. Please use `tf.keras.layers.Dense` instead.
  from ipykernel import kernelapp as app
/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel_launcher.py:16: UserWarning: `tf.layers.dense` is deprecated and will be removed in a future version. Please use `tf.keras.layers.Dense` instead.
  app.launch_new_instance()

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.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 ansioso. Colete estatísticas para as primeiras 3 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 verificar o caderno de validação de correção e equivalência numérica para obter conselhos adicionais para 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 de passagem única
  2. Validação de equivalência numérica de treinamento de modelo para algumas etapas
  3. Desempenho de inferência de benchmark
  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 .

Observe que é possível executar APIs de sessão e execução antecipada 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

    # abosolute 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. Verifique os documentos da API para obter detalhes.
  • Para execução antecipada, print e tf.print imprimem o valor do tensor. Mas print pode envolver a cópia do dispositivo para o host, o que pode tornar seu código mais lento.
  • Para o modo gráfico incluindo uso dentro de tf.function , você precisa usar tf.print para imprimir o valor real do tensor. tf.print é compilado em um op no gráfico, enquanto print e logging.info apenas registram 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 retorno de chamada para monitorar métricas e variáveis. Verifique como usar retornos de chamada personalizados com logs dict e 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]

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

tf.distribute.Strategy

  • Se a função tf.print tf.function executada nos trabalhadores, por exemplo, ao usar TPUStrategy ou ParameterServerStrategy , você precisará verificar os logs do servidor de trabalhadores/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 modelos de API Sequencial e Funcional, se você deseja imprimir valores, por exemplo, entradas de modelo ou recursos intermediários após algumas camadas, você tem as seguintes opções.
    1. Escreva uma camada personalizada que tf.print as entradas.
    2. Inclua as saídas intermediárias que você 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 ponto de verificação, em vez disso, escreva uma camada de subclasse personalizada. Verifique os documentos da API para obter mais detalhes.
  • Você não pode tf.print saídas intermediárias em um tf.keras.callbacks.LambdaCallback se você não tiver acesso aos valores reais, mas apenas aos objetos simbólicos do tensor Keras.

Opção 1: escreva 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])

[[-0.327884018]
 [-0.109294668]
 [-0.218589336]]
1/1 [==============================] - 0s 280ms/step - loss: 0.6077
<keras.callbacks.History at 0x7f63d46bf190>

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

Observe que, nesse caso, você pode 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 no terminal e no Colab para inspecionar valores intermediários para depuração.

Visualize gráfico com TensorBoard

Você pode examinar o gráfico do TensorFlow com o TensorBoard . O TensorBoard também é compatível com 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 saídas intermediárias entre o modelo TF1.x e o modelo TF2 migrado por meio do processo de treinamento e ver se os valores estão conforme o esperado.

Perfilador do TensorFlow

O TensorFlow Profiler pode ajudar você a visualizar o cronograma de execução em GPUs/TPUs. Você pode conferir esta demonstração do Colab para seu uso básico.

MultiProcessRunner

MultiProcessRunner é uma ferramenta útil ao depurar com MultiWorkerMirroredStrategy e ParameterServerStrategy. Você pode dar uma olhada neste exemplo concreto para seu uso.

Especificamente para os casos dessas duas estratégias, recomenda-se 1) não apenas ter testes de unidade para cobrir seu fluxo, 2) mas também tentar reproduzir falhas usando-o em teste de unidade para evitar lançar tarefas distribuídas reais toda vez que tentarem um conserto.