偵錯 TensorFlow 2 遷移的訓練管線

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視原始碼 下載筆記本

此筆記本示範如何在遷移至 TensorFlow 2 (TF2) 時偵錯訓練管線。它包含以下組件

  1. 偵錯訓練管線的建議步驟和程式碼範例
  2. 偵錯工具
  3. 其他相關資源

一個假設是您有 TensorFlow 1 (TF1.x) 程式碼和經過訓練的模型以進行比較,並且您想要建構一個 TF2 模型,以達到類似的驗證準確度。

此筆記本涵蓋偵錯訓練/推論速度或記憶體使用率的效能問題。

偵錯工作流程

以下是偵錯 TF2 訓練管線的一般工作流程。請注意,您不需要依序遵循這些步驟。您也可以使用二元搜尋方法,在中間步驟中測試模型,並縮小偵錯範圍。

  1. 修正編譯和執行階段錯誤

  2. 單次正向傳遞驗證(在另一個指南中)

    a. 在單一 CPU 裝置上

    • 驗證變數僅建立一次
    • 檢查變數計數、名稱和形狀是否相符
    • 重設所有變數,檢查在停用所有隨機性的情況下數值等效性
    • 對齊隨機數字產生,檢查推論中的數值等效性
    • (選用)檢查檢查點是否已正確載入,以及 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.*"

單次正向傳遞驗證

單次正向傳遞驗證,包括檢查點載入,在另一個 colab 中涵蓋。

import sys
import unittest
import numpy as np

import tensorflow as tf
import tensorflow.compat.v1 as v1
2024-01-17 02:21:07.536045: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory

模型訓練數值等效性驗證幾個步驟

設定模型組態並準備假資料集。

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 中使用相同的種子,

此工具提供兩種測試模式

  1. constant,它為每個單一操作使用相同的種子,無論它被呼叫了多少次,以及
  2. num_random_ops,它使用先前觀察到的有狀態隨機操作的數量作為操作種子。

這適用於用於建立和初始化變數的有狀態隨機操作,以及用於計算中的有狀態隨機操作(例如用於 dropout 層)。

random_tool = v1.keras.utils.DeterministicRandomTestTool(mode='num_random_ops')
WARNING:tensorflow:From /tmpfs/tmp/ipykernel_9596/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)
2024-01-17 02:21:10.121960: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2024-01-17 02:21:10.122074: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory
2024-01-17 02:21:10.122150: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublasLt.so.11'; dlerror: libcublasLt.so.11: cannot open shared object file: No such file or directory
2024-01-17 02:21:10.122222: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcufft.so.10'; dlerror: libcufft.so.10: cannot open shared object file: No such file or directory
2024-01-17 02:21:10.189341: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusparse.so.11'; dlerror: libcusparse.so.11: cannot open shared object file: No such file or directory
2024-01-17 02:21:10.189554: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1850] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://tensorflow.dev.org.tw/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
/tmpfs/tmp/ipykernel_9596/1984550333.py:14: UserWarning: `tf.layers.dense` is deprecated and will be removed in a future version. Please use `tf.keras.layers.Dense` instead.
  out_1 = tf.compat.v1.layers.dense(features, units=params['layer_1_size'])
/tmpfs/tmp/ipykernel_9596/1984550333.py:15: UserWarning: `tf.layers.dense` is deprecated and will be removed in a future version. Please use `tf.keras.layers.Dense` instead.
  out_2 = tf.compat.v1.layers.dense(out_1, units=params['layer_2_size'])
/tmpfs/tmp/ipykernel_9596/1984550333.py:16: UserWarning: `tf.layers.dense` is deprecated and will be removed in a future version. Please use `tf.keras.layers.Dense` instead.
  logits = tf.compat.v1.layers.dense(out_2, units=params['num_classes'])

定義 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 和 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 內使用,您需要使用 tf.print 來列印實際的張量數值。tf.print 會編譯到圖形中的 op,而 printlogging.info 僅在追蹤時記錄,這通常不是您想要的。
  • tf.print 也支援列印複合張量,例如 tf.RaggedTensortf.sparse.SparseTensor
  • 您也可以使用回呼來監控指標和變數。請查看如何將自訂回呼與 logs dictself.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]
Tensor("add:0", shape=(1,), dtype=float32)
[2]

tf.distribute.Strategy

  • 如果包含 tf.printtf.function 在 worker 上執行,例如當使用 TPUStrategyParameterServerStrategy 時,您需要檢查 worker/參數伺服器記錄以找到列印的值。
  • 對於 printlogging.info,當使用 ParameterServerStrategy 時,記錄將在協調器上列印,而當使用 TPU 時,記錄將在 worker0 上的 STDOUT 上列印。

tf.keras.Model

  • 當使用 Sequential 和 Functional API 模型時,如果您想要列印數值,例如模型輸入或某些層之後的中間特徵,您有以下選項。
    1. 編寫一個自訂層,使用 tf.print 列印輸入。
    2. 在模型輸出中包含您想要檢查的中間輸出。
  • tf.keras.layers.Lambda 層具有 (反)序列化限制。為了避免檢查點載入問題,請改為編寫自訂子類層。查看 API 文件以取得更多詳細資訊。
  • 如果您無法存取實際數值,而只能存取符號 Keras 張量物件,則無法在 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])
[[-0.327884018]
 [-0.109294683]
 [-0.218589365]]
1/1 [==============================] - 0s 273ms/step - loss: 0.6077
<keras.callbacks.History at 0x7effa3fcad30>

選項 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 上使用。TensorBoard 是一個很棒的可視化摘要工具。您可以使用它來比較 TF1.x 模型和遷移的 TF2 模型在訓練過程中的學習率、模型權重、梯度比例、訓練/驗證指標,甚至模型中間輸出,並查看數值是否如預期。

TensorFlow Profiler

TensorFlow Profiler 可以幫助您可視化 GPU/TPU 上的執行時間軸。您可以查看此 Colab 範例以了解其基本用法。

MultiProcessRunner

MultiProcessRunner 是一個有用的工具,當使用 MultiWorkerMirroredStrategy 和 ParameterServerStrategy 進行偵錯時。您可以查看這個具體範例以了解其用法。

特別是對於這兩種策略的情況,建議您 1) 不僅進行單元測試以涵蓋其流程,2) 而且還嘗試在單元測試中使用它來重現失敗,以避免每次嘗試修復時都啟動真正的分散式作業。