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.

Depurar una canalización de entrenamiento migrada de TensorFlow 2

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver código fuente en GitHub

Descargar bloc de notas

Este bloc de notas muestra cómo depurar una canalización de entrenamiento al migrar a TensorFlow 2 (TF2). Consta de los siguientes componentes:

1. Pasos sugeridos y muestras de código para depurar la canalización del entrenamiento

2. Herramientas para la depuración

3. Otros recursos relacionados

Una premisa es que usted tiene el código en TensorFlow 1 (TF1.x) y los modelos entrenados para su comparación, y desea construir un modelo TF2 que alcance una precisión de validación similar.

Este bloc de notas NO cubre la depuración de problemas de rendimiento para la velocidad de entrenamiento/inferencia o el uso de memoria.

Depuración del flujo de trabajo

Más abajo encontrará un flujo de trabajo general para depurar sus canalizaciones de entrenamiento en TF2. Tenga en cuenta que no es necesario que siga estos pasos en orden. También puede usar un enfoque de búsqueda binaria en el que pruebe el modelo en un paso intermedio y reduzca el alcance de la depuración.

1. Corregir errores de compilación y runtime

2. Validación de una sola pasada hacia delante (en una guía por separado)

   a. En un dispositivo con una sola CPU

   • Verificar que las variables se crean una sola vez

   • Compruebe que los recuentos, nombres y formas de las variables coinciden

   • Restablecer todas las variables, comprobar la equivalencia numérica con toda la aleatoriedad desactivada

   • Alinear la generación de números aleatorios, comprobar la equivalencia numérica en la inferencia

   • (Opcional) Los puntos de verificación se cargan correctamente y los modelos TF1.x/TF2 generan una salida idéntica

   b. En un único dispositivo GPU/TPU

   c. Con estrategias multidispositivo

3. Modelar el entrenamiento de validación de equivalencia numérica para unos pocos pasos (ejemplos de código disponibles a continuación)

   a. Validación de un solo paso de entrenamiento usando datos pequeños y fijos en un solo dispositivo con CPU. Específicamente, compruebe la equivalencia numérica para los siguientes componentes

   • computación de pérdidas

   • métricas

   • tasa de aprendizaje

   • cálculo y actualización del gradiente

   b. Compruebe las estadísticas después del entrenamiento de 3 o más pasos para verificar los comportamientos del optimizador como el impulso, aún con datos fijos en un solo dispositivo de CPU.

   c. En un único dispositivo GPU/TPU

   d. Con estrategias multidispositivo (consulte la introducción de MultiProcessRunner en la parte inferior)

4. Pruebas de convergencia de principio a fin en un conjunto de datos reales

   a. Comprobar los comportamientos de entrenamiento con TensorBoard

   • use optimizadores sencillos, por ejemplo SGD, y estrategias de distribución sencillas, por ejemplo tf.distribute.OneDeviceStrategy primero

   • métricas de entrenamiento

   • métricas de evaluación

   • averiguar cuál es la tolerancia razonable para la aleatoriedad inherente

   b. Comprobar la equivalencia con optimizador avanzado/programador de tasa de aprendizaje/estrategias de distribución

   c. Comprobar la equivalencia al usar precisión mixta

5. Puntos de referencia adicionales del producto

Preparación

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

Validación de una sola pasada hacia delante

La validación de una sola pasada, incluida la carga de puntos de verificación, se trata en otro colab.

import sys
import unittest
import numpy as np

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

Validación de la equivalencia numérica del entrenamiento del modelo por unos pocos pasos

Ajuste la configuración del modelo y prepare un conjunto de datos 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 el 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])

La siguiente clase v1.keras.utils.DeterministicRandomTestTool ofrece un administrador de contexto scope() que puede hacer que las operaciones aleatorias con estado usen la misma semilla en ambos grafos/sesiones TF1 y ejecución eager,

La herramienta ofrece dos modos de prueba:

1. constant que usa la misma semilla para cada operación sin importar cuántas veces haya sido llamada y,

2. num_random_ops que usa el número de operaciones aleatorias con estado observadas previamente como semilla de operación.

Esto se aplica tanto a las operaciones aleatorias con estado usadas para crear e inicializar variables, como a las operaciones aleatorias con estado usadas en el cálculo (como para las capas abandonadas).

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

Ejecute el modelo TF1.x en modo grafo. Recopile las estadísticas de los 3 primeros pasos de entrenamiento para comparar la equivalencia 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 el 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

Ejecute el modelo TF2 en modo eager. Recopile las estadísticas de los 3 primeros pasos de entrenamiento para comparar la equivalencia 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 la equivalencia numérica de los primeros pasos del entrenamiento.

También puede consultar el bloc de notas Validación de la corrección y equivalencia numérica si desea consejos adicionales para la equivalencia 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])

Pruebas de unidad

Existen algunos tipos de pruebas de unidad que pueden ayudar a depurar su código de migración.

1. Validación de una sola pasada hacia delante

2. Validación de la equivalencia numérica del entrenamiento del modelo por unos pocos pasos

3. Comparar los resultados del rendimiento de inferencia

4. El modelo entrenado realiza predicciones correctas sobre puntos de datos fijos y simples

Puede usar @parameterized.parameters para probar modelos con diferentes configuraciones. Detalles con muestra de código.

Tenga en cuenta que es posible ejecutar APIs de sesión y ejecución eager en el mismo caso de prueba. Los fragmentos de código a continuación muestran cómo.

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)

Herramientas de depuración

tf.print

tf.print vs print/logging.info

• Con argumentos configurables, tf.print puede mostrar recursivamente los primeros y últimos elementos de cada dimensión para los tensores impresos. Consulte la documentación de la API para más detalles.

• Para la ejecución eager, tanto print como tf.print imprimen el valor del tensor. Pero print puede suponer una copia de dispositivo a host, lo que potencialmente puede ralentizar su código.

• Para el modo gráfico, incluyendo el uso dentro de tf.function, necesita usar tf.print para imprimir el valor real del tensor. tf.print se compila en una op en el grafo, mientras que print y logging.info sólo hacen registros en el tiempo de seguimiento, que a menudo no es lo que usted desea.

• tf.print también permite imprimir tensores compuestos como tf.RaggedTensor y tf.sparse.SparseTensor.

• También puede usar una retrollamada para monitorear métricas y variables. Consulte cómo usar retrollamadas personalizadas con diccionarios de registros y 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

• Si el tf.function que contiene tf.print se ejecuta en los trabajadores, por ejemplo al usar TPUStrategy o ParameterServerStrategy, tiene que comprobar los registros del servidor de trabajadores/parámetros para encontrar los valores impresos.

• Para print o logging.info, los registros se imprimirán en el coordinador cuando se use ParameterServerStrategy, y los registros se imprimirán en STDOUT en worker0 cuando se usen TPUs.

tf.keras.Model

• Al usar modelos API secuenciales y funcionales, si desea imprimir valores, por ejemplo, entradas del modelo o características intermedias después de algunas capas, tiene las siguientes opciones.

  1. Escriba una capa personalizada que imprima por tf.print las entradas.

  2. Incluya las salidas intermedias que desee inspeccionar en las salidas del modelo.

• tf.keras.layers.Lambda las capas tienen limitaciones de (de)serialización. Para evitar problemas de carga de puntos de verificación, escriba en su lugar una capa personalizada de subclases. Consulte la documentación de API para más detalles.

• No puede imprimir por tf.print salidas intermedias en una retrollamada tf.keras.callbacks.LambdaCallback si no tiene acceso a los valores reales, sino sólo a los objetos tensores simbólicos de Keras.

Opción 1: escriba una capa 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])

Opción 2: incluya las salidas intermedias que desee inspeccionar en las salidas del modelo.

Tenga en cuenta que, en tal caso, puede necesitar algunas personalizaciones 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

Puede usar pdb tanto en el terminal como en Colab para inspeccionar los valores intermedios para depuración.

Visualizar grafo con TensorBoard

Puede examinar el grafo TensorFlow con TensorBoard. TensorBoard también está soportado en colab. TensorBoard es una gran herramienta para visualizar sumarios. Puede usarlo para comparar el ritmo de aprendizaje, las ponderaciones del modelo, la escala de gradiente, las métricas de entrenamiento/validación o incluso las salidas intermedias del modelo entre el modelo TF1.x y el modelo TF2 migrado a través del proceso de entrenamiento y ver si los valores son los esperados.

TensorFlow Profiler

TensorFlow Profiler puede ayudarle a visualizar la línea de tiempo de ejecución en GPUs/TPUs. Puede consultar esta Demo de Colab para ver su uso básico.

MultiProcessRunner

MultiProcessRunner es una herramienta útil a la hora de depurar con MultiWorkerMirroredStrategy y ParameterServerStrategy. Puede echar un vistazo a este ejemplo concreto para ver su uso.

Específicamente para los casos de estas dos estrategias, se recomienda 1) no sólo tener pruebas de unidad para cubrir su flujo, 2) sino también intentar reproducir fallos usándolo en pruebas de unidad para evitar lanzar un trabajo real distribuido cada vez que se intente una reparación.