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.

TF2 に移行されたトレーニングパイプラインをデバッグする

TensorFlow.org で表示

Google Colab で実行

GitHub でソースを表示

ノートブックをダウンロード

このノートブックでは、TensorFlow 2（TF2）に移行されたトレーニングパイプラインをデバッグする方法を説明します。内容は以下のとおりです。

1. トレーニングパイプラインをデバッグするための推奨される手順とコードサンプル

2. デバッグ用ツール

3. その他の関連リソース

比較用の TensorFlow 1（TF1.x）コードとトレーニング済みモデルがあり、同等の検証精度を達成する TF2 モデルを構築することを前提とします。

このノートブックは、トレーニングや推論の速度やメモリ使用量に関するデバッグパフォーマンスの問題は取り上げません。

デバッグワークフロー

以下は、TF2 トレーニングパイプラインをデバッグするための一般的なワークフローです。これらの手順を順番に実行する必要はありません。中間ステップでモデルをテストし、デバッグ範囲を絞り込む二分探索アプローチを使用することもできます。

1. コンパイルエラーとランタイムエラーを修正する

2. シングルフォワードパスの検証（別のガイド）

   a. 単一の CPU デバイスの場合

   • 変数が 1 回だけ作成されることを確認する

   • 変数の数、名前、形状が一致していることを確認する

   • すべての変数をリセットし、すべてのランダム性を無効にして数値の等価性をチェックする

   • 乱数生成の調整、推論における数値的等価性をチェックする

   • （オプション）チェックポイントが正しく読み込まれ、TF1.x/TF2 モデルが同一の出力を生成することを確認する

   b. 単一の GPU/TPU デバイスの場合

   c. マルチデバイスストラテジー

3. 数ステップのモデルトレーニングの数値的等価性の検証（コードサンプルは以下で入手可能）

   a. 単一の CPU デバイスでの小規模な固定データを使用した単一のトレーニングステップの検証。具体的には、次のコンポーネントの数値的等価性を確認する

   • 損失計算

   • 指標

   • 学習率

   • 勾配の計算と更新

   b. 3 つ以上のステップをトレーニングした後に統計をチェックして、モメンタムなどのオプティマイザの動作を検証する。単一の CPU デバイスで固定データを使用する。

   c. 単一の GPU/TPU デバイス

   d. マルチデバイスストラテジーを使用（以下の MultiProcessRunner のイントロを参照）

4. 実際のデータセットでのエンドツーエンドの収束テスト

   a. TensorBoard でトレーニングの動作を確認する

   • 単純なオプティマイザを使用する（SGD と単純な分布戦略。 最初に tf.distribute.OneDeviceStrategy を使用する）。

   • トレーニング指標

   • 評価指標

   • 固有のランダム性に対する妥当な許容範囲を把握する

   b. 高度なオプティマイザ/学習率スケジューラ/分散ストラテジーとの同等性をチェックする

   c. 混合精度使用時の同等性をチェックする

5. 追加の乗積ベンチマーク

セットアップ

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

1 つのフォワードパスの検証

チェックポイントの読み込みを含む 1 つのフォワードパスの検証については、別の 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 グラフ/セッションと Eager execution の両方でステートフルなランダム演算が同じシードを使用できるようになります。

このツールには、次の 2 つのテストモードがあります。

1. constant は、呼び出された回数に関係なく、1 つの演算ごとに同じシードを使用します。

2. num_random_ops は、以前に観測されたステートフルなランダム演算の数を演算シードとして使用します。

これは、変数の作成と初期化に使用されるステートフルなランダム演算と、計算で使用されるステートフルなランダム演算（ドロップアウトレイヤーなど）の両方に適用されます。

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

TF1.x モデルを Graph モードで実行します。数値的等価性を比較するために、最初の 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)

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

TF2 モデルを eager モードで実行します。数値的等価性を比較するために、最初の 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. 1 つのフォワードパスの検証

2. 数ステップのモデルトレーニングの数値的等価性検証

3. ベンチマーク推論性能

4. トレーニング済みのモデルが固定された単純なデータポイントに対して正しい予測を行う

@parameterized.parameters を使用して、さまざまな構成でモデルをテストできます。詳細（コードサンプル付き）はこちらを参照してください。

セッション API と Eager execution を同じテストケースで実行できます。以下のコードスニペットは、その方法を示しています。

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)

デバッグツール

tf.print

tf.print と print/logging.info の比較

• 構成可能な引数を使用すると、tf.print は出力されたテンソルの各次元の最初と最後のいくつかの要素を再帰的に表示できます。詳細については、API ドキュメントを参照してください。

• Eager execution では、print と tf.print の両方がテンソルの値を出力します。ただし、print にはデバイスからホストへのコピーが含まれる場合があり、コードが遅くなる可能性があります。

• tf.function 内での使用を含む Graph モードでは、tf.print を使用して実際のテンソル値を出力する必要があります。tf.print はグラフ内の演算にコンパイルされますが、print と logging.info はトレース時にしかログに記録されません（多くの場合、これは希望されないことだと思います）。

• tf.print は、tf.RaggedTensor や tf.sparse.SparseTensor などの複合テンソルの出力もサポートしています。

• また、コールバックを使用して、指標と変数を監視することもできます。logs dict と self.model 属性でカスタムコールバックを使用する方法を確認してください。

tf.print と tf.function 内の print の比較

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

• tf.print を含む tf.function がワーカーで実行される場合（たとえば、TPUStrategy または ParameterServerStrategy を使用する場合）、出力された値を見つけるには、ワーカー/パラメータサーバーログを確認する必要があります。

• print または logging.info の場合、ParameterServerStrategy を使用するとログがコーディネータに出力され、TPU を使用する場合は、ログは worker0 の STDOUT に出力されます。

tf.keras.Model

• Sequential API モデルと Functional API モデルを使用する場合、モデル入力やいくつかのレイヤーの後の中間特徴などの値を出力する場合は、次のオプションがあります。

  1. 入力を tf.print するカスタムレイヤーを作成します。

  2. 調査する中間出力をモデル出力に含めます。

• tf.keras.layers.Lambda レイヤーには（逆）シリアル化の制限があります。チェックポイントの読み込みの問題を回避するには、カスタムサブクラス化されたレイヤーを記述します。詳しくは、API ドキュメントを参照してください。

• 実際の値にアクセスできない場合、tf.keras.callbacks.LambdaCallback で中間出力を tf.print することはできませんが、シンボリック 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])

オプション 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 の両方で pdb を使用して、デバッグ用の中間値を調べることができます。

TensorBoard でグラフを可視化する

TensorBoard を使用すると TensorFlow のグラフを調べられます。TensorBoard は、要約を視覚化する優れたツールで colab でもサポートされています。これを使用して、トレーニングプロセスを通じて TF1.x モデルと移行された TF2 モデルの間で学習率、モデルの重み、勾配スケーリング、トレーニング/検証指標、および、モデルの中間出力を比較し、値が期待どおりになっているかどうかを確認できます。

TensorFlow Profiler

TensorFlow Profiler は、GPU/TPU での実行タイムラインを視覚化するのに役立ちます。基本的な使い方については、この Colab デモを参照してください。

MultiProcessRunner

MultiProcessRunner は、MultiWorkerMirroredStrategy と ParameterServerStrategy でデバッグする際に便利なツールです。使用法については、この具体的な例を参照してください。

特にこれら 2 つのストラテジーのケースでは、1) フローをカバーする単体テストを用意し、2) 単体テストでこれを使用して失敗を再現してみることをお勧めします。これは、修正を試みるたびに実際の分散ジョブが起動されることを避けるためです。