Отладка конвейера обучения, перенесенного в TF2

Посмотреть на TensorFlow.org

Запустить в Google Colab

Посмотреть исходный код на GitHub

Скачать блокнот

В этой записной книжке показано, как отлаживать конвейер обучения при переходе на TF2. Он состоит из следующих компонентов:

1. Предлагаемые шаги и примеры кода для отладки конвейера обучения
2. Инструменты для отладки
3. Другие связанные ресурсы

Одно из предположений: у вас есть код TF1.x и обученные модели для сравнения, и вы хотите построить модель TF2, обеспечивающую аналогичную точность проверки.

В этой записной книжке НЕ рассматриваются проблемы производительности отладки для скорости обучения/вывода или использования памяти.

Рабочий процесс отладки

Ниже приведен общий рабочий процесс для отладки конвейеров обучения TF2. Обратите внимание, что вам не нужно выполнять эти шаги по порядку. Вы также можете использовать подход бинарного поиска, при котором вы тестируете модель на промежуточном этапе и сужаете область отладки.

1. Исправление ошибок компиляции и выполнения

2. Проверка единого прохода вперед (в отдельном руководстве )

   а. На устройстве с одним процессором

   • Убедитесь, что переменные создаются только один раз
   • Проверьте количество переменных, имена и формы совпадают
   • Сбросить все переменные, проверить числовую эквивалентность с отключенной случайностью
   • Выровняйте генерацию случайных чисел, проверьте числовую эквивалентность в выводе
   • (Необязательно) Проверьте, правильно ли загружены контрольные точки, и модели TF1.x/TF2 генерируют идентичный вывод.

   б. На одном устройстве GPU/TPU

   в. Стратегии для нескольких устройств

3. Проверка числовой эквивалентности обучения модели для нескольких шагов (примеры кода доступны ниже)

   а. Проверка одного шага обучения с использованием небольших и фиксированных данных на устройстве с одним процессором. В частности, проверьте числовую эквивалентность для следующих компонентов.

   • расчет убытков
   • показатели
   • скорость обучения
   • вычисление и обновление градиента

   б. Проверка статистики после обучения 3 или более шагов для проверки поведения оптимизатора, такого как импульс, по-прежнему с фиксированными данными на устройстве с одним ЦП

   в. На одном устройстве GPU/TPU

   д. Со стратегиями для нескольких устройств (проверьте введение MultiProcessRunner внизу)

4. Сквозное тестирование покрытия на реальном наборе данных

   а. Проверяйте тренировочное поведение с помощью TensorBoard

   • сначала используйте простые оптимизаторы, например SGD, и простые стратегии распространения, например tf.distribute.OneDeviceStrategy
   • показатели обучения
   • показатели оценки
   • выяснить, какова разумная терпимость к присущей случайности

   б. Проверьте эквивалентность с помощью расширенного оптимизатора/планировщика скорости обучения/стратегий распределения

   в. Проверяйте эквивалентность при использовании смешанной точности

5. Дополнительные тесты продукта

Настраивать

pip uninstall -y -q tensorflow

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

Проверка однократного прямого прохода

Проверка одиночного прямого прохода, включая загрузку контрольной точки, рассматривается в другом colab .

import sys
import unittest
import numpy as np

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

Проверка числовой эквивалентности обучения модели за несколько шагов

Настройте конфигурацию модели и подготовьте поддельный набор данных.

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

Определите модель 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])

Следующий класс v1.keras.utils.DeterministicRandomTestTool предоставляет диспетчер контекста scope() , который может заставить случайные операции с отслеживанием состояния использовать одно и то же начальное значение как для графов/сеансов TF1, так и для нетерпеливого выполнения.

Инструмент предоставляет два режима тестирования:

1. constant , которая использует одно и то же семя для каждой отдельной операции, независимо от того, сколько раз она была вызвана и,
2. num_random_ops , который использует количество ранее наблюдаемых случайных операций с отслеживанием состояния в качестве начального значения операции.

Это относится как к случайным операциям с отслеживанием состояния, используемым для создания и инициализации переменных, так и к случайным операциям с отслеживанием состояния, используемым в вычислениях (например, для выпадающих слоев).

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.

Запустите модель TF1.x в графическом режиме. Соберите статистику для первых 3 шагов обучения для сравнения числовой эквивалентности.

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

Определите модель 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

Запустите модель TF2 в активном режиме. Соберите статистику для первых 3 шагов обучения для сравнения числовой эквивалентности.

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

Сравните числовую эквивалентность для первых нескольких шагов обучения.

Вы также можете проверить блокнот Проверка правильности и числовой эквивалентности для получения дополнительных советов по числовой эквивалентности.

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

Модульные тесты

Существует несколько типов модульного тестирования, которые могут помочь в отладке кода миграции.

1. Проверка однократного прямого прохода
2. Проверка числовой эквивалентности обучения модели за несколько шагов
3. Сравнительная производительность логического вывода
4. Обученная модель делает правильные прогнозы на фиксированных и простых точках данных.

Вы можете использовать @parameterized.parameters для тестирования моделей с различными конфигурациями. Подробности с примером кода .

Обратите внимание, что в одном и том же тестовом примере можно запускать сеансовые API и активное выполнение. Фрагменты кода ниже показывают, как это сделать.

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)

Инструменты отладки

tf.print

tf.print против print/logging.info

• С настраиваемыми аргументами tf.print может рекурсивно отображать несколько первых и последних элементов каждого измерения для печатных тензоров. Подробности смотрите в документации по API .
• Для быстрого выполнения как print , так и tf.print печатают значение тензора. Но print может включать копирование с устройства на хост, что потенциально может замедлить работу вашего кода.
• Для режима графика, включая использование внутри tf.function , вам нужно использовать tf.print для печати фактического значения тензора. tf.print компилируется в операцию на графике, тогда как print и logging.info регистрируются только во время трассировки, что часто не то, что вам нужно.
• tf.print также поддерживает печать составных тензоров, таких как tf.RaggedTensor и tf.sparse.SparseTensor .
• Вы также можете использовать обратный вызов для мониторинга метрик и переменных. Пожалуйста, проверьте, как использовать пользовательские обратные вызовы с атрибутами logs dict и self.model .

tf.print против печати внутри 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.Стратегия

• Если tf.function , содержащий tf.print , выполняется на рабочих процессах, например, при использовании TPUStrategy или ParameterServerStrategy , вам необходимо проверить журналы рабочего процесса/сервера параметров, чтобы найти напечатанные значения.
• Для print или logging.info журналы будут распечатаны на координаторе при использовании ParameterServerStrategy , а журналы будут распечатаны на STDOUT на worker0 при использовании TPU.

tf.keras.Модель

• При использовании моделей Sequential и Functional API, если вы хотите распечатать значения, например, входные данные модели или промежуточные объекты после некоторых слоев, у вас есть следующие варианты.
  1. Напишите пользовательский слой , который tf.print вводит.
  2. Включите промежуточные выходные данные, которые вы хотите проверить, в выходные данные модели.
• Слои tf.keras.layers.Lambda имеют ограничения (де)сериализации. Чтобы избежать проблем с загрузкой контрольных точек, вместо этого напишите пользовательский слой с подклассами. Дополнительные сведения см. в документации по API .
• Вы не можете tf.print промежуточные выходные данные в tf.keras.callbacks.LambdaCallback , если у вас нет доступа к фактическим значениям, а только к символическим объектам тензора Keras.

Вариант 1: написать пользовательский слой

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>

Вариант 2: включите промежуточные выходные данные, которые вы хотите проверить, в выходные данные модели.

Обратите внимание, что в таком случае вам могут потребоваться некоторые настройки для использования 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 как в терминале, так и в Colab для проверки промежуточных значений для отладки.

Визуализируйте график с помощью TensorBoard

Вы можете изучить граф TensorFlow с помощью TensorBoard . TensorBoard также поддерживается в colab . TensorBoard — отличный инструмент для визуализации сводок. Вы можете использовать его для сравнения скорости обучения, веса модели, шкалы градиента, показателей обучения/проверки или даже для моделирования промежуточных результатов между моделью TF1.x и перенесенной моделью TF2 в процессе обучения и проверки того, выглядят ли значения так, как ожидалось.

Профилировщик TensorFlow

TensorFlow Profiler может помочь вам визуализировать временную шкалу выполнения на GPU/TPU. Вы можете проверить эту демонстрацию Colab для ее основного использования.

MultiProcessRunner

MultiProcessRunner — полезный инструмент при отладке с помощью MultiWorkerMirroredStrategy и ParameterServerStrategy. Вы можете взглянуть на этот конкретный пример его использования.

В частности, для случаев этих двух стратегий вам рекомендуется: 1) не только иметь модульные тесты для покрытия их потока, 2) но также попытаться воспроизвести сбои, используя их в модульном тесте, чтобы избежать запуска реального распределенного задания каждый раз, когда они пытаются исправление.