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에서 소스 보기

노트북 다운로드하기

이 노트북은 TF2로 마이그레이션할 때 훈련 파이프라인을 디버그하는 방식을 설명하며 다음 구성 요소로 구성되어 있습니다.

1. 훈련 파이프라인 디버깅 권장 단계 및 코드 샘플

2. 디버깅 도구

3. 기타 관련 리소스

비교에 사용하는 TF1.x 코드와 훈련된 모델이 있고 유사한 검증 정확성을 달성하는 TF2 모델을 구축하려 한다고 가정합니다.

이 노트북은 훈련/추론 속도 또는 메모리 사용량에 대한 디버깅 성능 문제를 다루지 않습니다.

워크플로 디버깅하기

다음은 TF2 훈련 파이프라인을 디버깅하는 일반적인 워크플로입니다. 이 단계를 순서대로 따를 필요는 없습니다. 중간 단계에서 모델을 테스트하고 디버깅 범위를 좁히는 이진 검색 접근 방식을 사용할 수도 있습니다.

1. 컴파일 및 런타임 오류 수정

2. 단일 순방향 전달 검증(별도의 가이드)

   a. 단일 CPU 기기에서

   • 변수가 한 번만 생성되었는지 확인

   • 변수 개수, 이름 및 모양이 일치하는지 확인

   • 모든 변수 재설정, 모든 임의성이 비활성화된 상태에서 수치적 동등성 확인

   • 난수 생성 정렬, 추론에서 수치적 동등성 확인

   • (선택 사항)체크포인트가 제대로 로드되고 TF1.x/TF2 모델이 동일한 출력을 생성하는지 확인

   b. 단일 GPU/TPU 기기에서

   c. 멀티 기기 전략을 사용하는 경우

3. 몇 가지 단계의 모델 훈련 수치적 동등성 검증(아래의 코드 샘플 사용 가능)

   a. 단일 CPU 기기에서 고정된 작은 데이터를 사용하는 단일 훈련 단계 검증. 특히 다음 구성 요소에 대한 수치 동등성 확인

   • 손실 계산

   • 메트릭

   • 학습률

   • 그래디언트 계산 및 업데이트

   b. 단일 CPU 기기에 고정 데이터가 있는 모멘텀과 같은 옵티마이저 동작을 확인하기 위해 3개 이상의 단계를 훈련한 후 통계를 확인

   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.*"

단일 순방향 전달 검증

체크포인트 로드 등 단일 순방향 전달 검증은 다른 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 클래스는 상태 저장 임의 작업이 TF1 그래프/세션과 즉시 실행 모두에서 동일한 시드를 사용하도록 할 수 있는 컨텍스트 관리자 scope()를 제공합니다.

이 도구는 두 가지 테스트 모드를 제공합니다.

1. 호출된 횟수에 관계없이 모든 단일 연산에 대해 동일한 시드를 사용하는 constant

2. 이전에 관찰된 상태 저장 임의 연산 수를 연산 시드로 사용하는 num_random_ops

이는 변수 생성 및 초기화에 사용되는 상태 저장 임의 작업과 계산에 사용되는 상태 저장 임의 작업(예: 드롭아웃 레이어)에 모두 적용됩니다.

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

그래프 모드로 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)

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

Eager 모드로 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

    # 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 문서를 확인하세요.

• 즉시 실행의 경우 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]

tf.distribute.Strategy

• TPUStrategy 또는 ParameterServerStrategy를 사용하는 경우와 같이 tf.print를 포함하는 tf.function을 작업자에서 실행하는 경우 인쇄된 값을 찾으려면 작업자/매개변수 서버 로그를 확인해야 합니다.

• print 또는 logging.info의 경우, ParameterServerStrategy를 사용하면 로그가 코디네이터에 인쇄되고, TPU를 사용하면 로그가 worker0의 STDOUT에 인쇄됩니다.

tf.keras.Model

• 순차형 및 함수형 API 모델을 사용할 때 일부 레이어 뒤에 모델 입력 또는 중간 특성과 같은 값을 다음 옵션으로 인쇄할 수 있습니다.

  1. 입력을 tf.print하는 사용자 정의 레이어를 작성합니다.

  2. 검사하려는 중간 출력을 모델 출력에 포함합니다.

• tf.keras.layers.Lambda 레이어에는 (역)직렬화 제한이 있습니다. 체크포인트 로드 문제를 피하려면 대신 사용자 정의 서브 클래스된 레이어를 작성하세요. 자세한 내용은 API 문서를 확인하세요.

• 실제 값에 액세스할 수 없는 경우 tf.keras.callbacks.LambdaCallback에서 중간 출력을 tf.print할 수 없고 대신 기호 케라스 텐서 객체만 사용할 수 있습니다.

옵션 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을 사용하려면 몇 가지 [맞춤 설정](https://www.tensorflow.org/guide/keras/customizing_what_happens_in_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에서도 지원됩니다. TensorBoard는 요약을 시각화하는 훌륭한 도구입니다. 이를 사용하여 학습률, 모델 가중치, 그래디언트 스케일, 훈련/검증 메트릭을 비교하거나 학습 프로세스를 통해 TF1.x 모델과 마이그레이션된 TF2 모델 사이의 중간 출력을 모델링하고 값이 예상대로 나타나는지 확인할 수 있습니다.

TensorFlow 프로파일러

TensorFlow Profiler를 사용하면 GPU/TPU의 실행 타임라인을 시각화할 수 있습니다. 이 Colab 데모에서 기본적인 사용법을 확인할 수 있습니다.

MultiProcessRunner

MultiProcessRunner는 MultiWorkerMirroredStrategy와 ParameterServerStrategy로 디버깅할 때 유용한 도구입니다. 사용법은 이 구체적인 예제를 참조하세요.

특히 이 두 가지 전략의 경우 1)플로를 처리하는 단위 테스트를 수행할 뿐만 아니라 2)수정을 시도할 때마다 실제 분산 작업을 시작하지 않도록 단위 테스트에서 이를 사용하여 실패를 재현하는 것이 좋습니다.