如何在不失去理智的情况下调试 TensorFlow 训练程序

一、说明

        关于tensorflow的调试，是一个难啃的骨头，除了要有耐力，还需要方法；本文假设您是一个很有耐力的开发者，为您提供一些方法；这些方法也许不容易驾驭，但是依然强调您只要有耐力，没有不能攻克的难题。

二、关于tensorflow调试的说法 

If debugging is the process of removing software bugs, then programming must be the process of putting them in.
Edsger Dijkstra. From https://www.azquotes.com/quote/561997

        在我之前的一些帖子（这里，这里和这里）中，我告诉了你一些关于我在Mobileye的团队（正式名称为Mobileye，英特尔公司）如何使用TensorFlow，Amazon SageMaker和Amazon s3来训练我们的深度神经网络在大量数据上。在这篇文章中，我想谈谈TensorFlow中的调试。

        众所周知，程序调试是软件开发的一个组成部分，并且花费在调试上的时间通常会超过编写原始程序所需的时间。

        调试是困难的，关于如何设计和实现自己的程序以增加错误的可重复性并简化根本原因分析的过程已经写了很多。

        在机器学习中，由于机器学习算法固有的随机性，以及算法通常在远程计算机上的专用硬件加速器上运行这一事实，调试任务变得复杂。

        由于使用了符号执行（又名图模式），TensorFlow 中的调试变得更加复杂，这提高了训练会话的运行时性能，但同时限制了自由读取图中任意张量的能力，而这种能力对于调试很重要。

        在这篇文章中，我将扩展调试TensorFlow训练程序的困难，并提供一些如何解决这些困难的建议。

        出于法律目的，我想澄清一下，尽管我精心挑选了副标题，但我不保证我在这里写的任何内容都会阻止你气馁。相反，我认为我几乎可以保证，尽管我写了什么，但在调试TensorFlow程序时，你可能会气馁而失去理智。 

        在开始之前，让我们澄清一下讨论的范围。

三、调试的类型

        在这篇文章的上下文中，调试是指识别代码或数据中的错误，导致培训课程突然崩溃的艺术。

        另一种调试，超出了本文的范围，指的是修复或调整未收敛的模型的任务，或者对某一类输入产生不令人满意的预测（例如，无法识别粉红色汽车的车辆检测模型）。此过程可能涉及定义和评估模型指标、收集和统计分析模型工件（例如梯度、激活和权重）、使用 TensorBoard 和 Amazon Sagemaker 调试器等工具、超参数调优、重新架构或使用增强和提升等技术修改数据输入。调整模型可能是一项极具挑战性、耗时且经常令人沮丧的任务。

3.1 错误的类型

在解决代码或数据中的错误领域，我喜欢区分两类错误：错误和怪物错误。

通过错误，我指的是相对容易重现的问题。错误的示例包括假设输入张量的大小与训练数据不匹配的模型，尝试连接不匹配的张量，或对无效数据类型执行 tf 操作。这些通常不依赖于特定的模型状态和数据，并且通常相对容易重现。它们不一定容易修复，但与怪物虫子相比，它们是儿戏。

怪物虫是偶尔发生和不可预测的虫子。仅在模型的特定状态、特定数据样本或模型状态和数据输入的特定组合上重现的 bug 可能会带来严重的挑战，并可能构成怪物 bug。

下面是一个基于真实事件的场景示例，它肯定会增加您的血压：

现在是星期五下午，您的模型已经成功训练了几天。损失似乎正在很好地收敛，你开始想象一个放松的，发布后的周末假期，在你选择的度假地点。你回头看了一眼屏幕，注意到，突然之间，没有任何警告，你的损失变成了NaN。“当然”，你心想，“这一定是由于一些完全随机的、瞬间的、宏观的故障”，然后你立即从上一个有效的模型检查点恢复训练。又过了几个小时，它又发生了，然后又发生了。现在你开始恐慌，周末天堂的梦幻画面现在被需要解决怪物虫子的诱人努力的想法所取代。

我们一会儿再来看这个令人悲伤的例子。但首先，让我们勾选一些强制性的“调试”复选框。

四、在 Tensorflow 中调试的提示

很多笔墨已经洒在调试的艺术上，更重要的是，开发可调试代码的艺术。在本节中，我将提到一些技术，因为它们与TensorFlow应用程序有关。这份清单绝不是全面的。

4.1 保存模型检查点

这可能是我将在这篇文章中写的最重要的事情。始终配置训练会话，使其定期保存模型的快照。

编程错误并不是您的培训可能崩溃的唯一原因......如果您在云中运行，则可能会遇到 Spot 实例终止或遇到内部服务器错误。如果在本地运行，则可能会断电，或者 GPU 可能会爆炸。如果您已经训练了几天，没有存储中间检查点，那么损害可能是极端的。如果您每小时保存一个检查点，那么您最多只丢失一个小时。TensorFlow提供了用于存储检查点的实用程序，例如keras模型检查点回调。您需要做的就是通过权衡存储检查点的开销与培训课程中计划外故障的成本来决定捕获此类快照的频率。

4.2 接触者追踪

        我向我的 Covid19 同时代人道歉，因为我为这个小节选择了标题，我只是无法抗拒。通过接触者追踪，我指的是跟踪输入到训练管道中的训练数据的能力。

        假设您的训练数据在 tfrecord 文件中分为 100,000 个，并且其中一个文件存在导致程序崩溃或停止的格式错误。缩小对有问题文件的搜索范围的一种方法是记录输入到管道中的每个文件。一旦你遇到崩溃，你可以回顾你的日志，看看最近输入的文件是什么。正如我在之前的文章中提到的，我们使用 Amazon SageMaker 管道模式功能进行训练。管道模式最近新增的管道模式服务器端日志记录了输入到管道中的文件。

记录进入管道的数据可以帮助一个人重现错误的能力，这将我们带到下一点。

4.3 错误再现 

        重现错误的难易程度直接影响解决错误的难易程度。我们总是希望编写代码以确保可重复性。这在TensorFlow程序中并不容易。机器学习应用通常包括对随机变量使用的依赖。我们随机初始化模型权重，随机增加数据，随机分片数据以进行分布式训练，随机应用 dropout，在每个 epoch 之前对输入数据进行洗牌，然后在创建批处理之前再次对其进行洗牌（使用 tf.dataset.shuffle）。我们可以用我们记录的伪随机种子来播种所有的伪随机操作，但请记住，可能有许多不同的地方引入了随机化，跟踪所有这些很容易成为记账的噩梦。我无法告诉你有多少次我以为我已经删除了随机化的所有元素，却发现我错过了一个。此外，还有一些随机过程无法设定种子。如果您使用多个进程来导入训练数据，则可能无法控制数据记录的实际馈送顺序（例如，如果在 tf.data.Options（） 中将experimental_deterministic设置为 false）。当然，您可以在每个样本进入管道时记录它们，但这会带来陡峭且可能令人望而却步的开销。

        最重要的是，虽然构建可重复的训练程序绝对是可能的，但我认为更明智的做法是接受非确定性，接受训练的不可重复性，并找到克服这种调试限制的方法。

4.4 模块化编程

        创建可调试程序的一个关键技术是以模块化方式构建应用程序。应用于 TensorFlow 训练循环，这意味着能够分别测试训练管道的不同子集，例如数据集、损失函数、不同的模型层和回调。这并不总是那么容易做到，因为一些训练模块（如损失函数）非常依赖于其他模块。但是有很大的创造力空间。例如，只需迭代数据集，同时应用数据集操作的子集，就可以在输入管道上测试不同的函数。可以通过创建仅运行损失函数或回调的应用程序来测试损失函数或回调。人们可以通过用虚拟损失函数代替损失函数来中和损失函数。我喜欢用多个输出点构建我的模型，即能够轻松修改模型中的层数，以测试不同层的影响。

        在构建程序时，您对程序的模块化和可调试性投入的考虑越多，以后遭受的痛苦就越少。

4.5 急切的执行力

        如果您是常规的 TensorFlow 用户，您可能遇到过“渴望执行模式”、“图形模式”和“tf 函数限定符”等术语。您可能听说过一些（有些误导性）语句，例如“在渴望执行模式下调试是小菜一碟”或“Tensorflow 2 在渴望执行模式下运行”。你可能像我一样，热切地投入到tensorflow源代码中，试图理解不同的执行模式，结果却在抽泣中崩溃了，你的自尊心终生破碎。为了全面了解它是如何工作的，我建议你阅读TensorFlow文档，并祝你好运。在这里，我们将只提及它与调试相关的要点。运行 TensorFlow 训练的最佳方法是在图形模式下运行它。图模式是一种符号执行模式，这意味着我们不能任意访问图张量。使用 tf.function 限定符包装的函数将在图形模式下运行。使用 tf.keras.model.fit 进行训练时，默认情况下，训练步骤在图形模式下执行。当然，无法访问任意图张量使得在图模式下进行调试变得困难。在预先执行模式下，您可以访问任意张量，甚至可以使用调试器进行调试（前提是您将断点放在 model.call（） 函数中的适当位置）。当然，当您在渴望执行模式下运行时，您的训练会运行得更慢。若要对模型进行编程以在预先执行模式下进行训练，需要调用 model.compile（） 函数，并将 run_eagerly 标志设置为 true。

        底线是，当你在训练时，在图形模式下运行，当你在调试时，在预先执行模式下运行。不幸的是，某些错误仅在图形模式下而不是在急切执行模式下重现的情况并不少见，这真是令人沮丧。此外，在本地环境中调试时，预先执行很有帮助，而在云中则不那么有用。它在调试怪物错误时通常不是很有用......除非您首先找到一种在本地环境中重现该错误的方法（下面将对此进行详细介绍）。

4.6 TensorFlow Logging和Debugging Utility

        尝试充分利用TensorFlow记录器。调试问题时，请将记录器设置为信息量最大的级别。

tf.debugging 模块提供了一堆断言实用程序以及数字检查函数。特别是，tf.debugging.enable_check_numerics实用程序有助于查明有问题的函数。

tf.print 函数可以打印出任意图形张量，这是一个额外的实用程序，我发现它对调试非常有用。

        最后但并非最不重要的一点是，添加您自己的打印日志（在代码的非图形部分中），以便更好地了解程序故障的位置。

4.7 解密 TensorFlow 错误消息

        有时，你会很幸运地收到TensorFlow错误消息。不幸的是，并不总是立即清楚如何使用它们。我经常收到同事的电子邮件，上面有神秘的TensorFlow消息，乞求帮助。当我看到消息时，例如：

tensorflow.python.framework.errors_impl.InvalidArgumentError: ConcatOp : Dimensions of inputs should match: shape[0] = [5,229376] vs. shape[2] = [3,1]

或

node DatasetToGraphV2 (defined at main.py:152) (1) Failed precondition: Failed to serialize the input pipeline graph: Conversion to GraphDef is not supported.

或

alueError: slice index -1 of dimension 0 out of bounds. for 'loss/strided_slice' (op: 'StridedSlice') with input shapes: [0], [1], [1], [1] and with computed input tensors: input[1] = <-1>, input[2] = <0>, input[3] = <1>.

        我问自己（稍微修改一下，使帖子对孩子友好）“我应该用它发出什么嘟嘟声？”或者“为什么友好的TensorFlow工程师不能给我更多的东西来工作？”。但我很快就让自己平静下来（有时在酒精饮料的帮助下），并说：“Chaim，别再被宠坏了。回去工作，感谢你得到了任何信息。您应该做的第一件事是尝试在预先执行模式下和/或使用调试器重现错误。不幸的是，如上所述，这并不总是有帮助。

        毫无疑问，上述消息不是很有帮助。但不要绝望。有时，在一些调查工作的帮助下，你会发现可能引导你走向正确方向的线索。浏览调用堆栈以查看它是否提供了任何提示。如果消息包含形状大小，请尝试将这些大小与图形中可能具有相同形状的张量进行匹配。当然，还可以在线搜索，看看其他人是否遇到过类似的问题以及在什么情况下遇到过。不要绝望。

4.8 在本地环境中运行

        当然，在本地环境中调试比在远程计算机或云中调试更容易。首次创建模型时尤其如此。在开始远程培训之前，您的目标应该是在当地环境中解决尽可能多的问题。否则，您最终可能会浪费大量时间和金钱。

        为了提高可重现性，应尝试使本地环境尽可能与远程环境相似。如果您在远程环境中使用 docker 映像或虚拟环境，请尝试在本地使用相同的映像或虚拟环境。（如果您的远程培训在 Amazon SageMaker 上进行，则可以拉取所使用的 docker 映像。

        当然，远程培训环境中的某些元素可能无法在本地复制。例如，您可能遇到了仅在使用 Amazon SageMaker 管道模式时重现的错误，该模式目前仅在云中运行时受支持。（在这种情况下，您可以考虑从 s3 访问数据的替代方法。

        我希望我能告诉你，这里描述的技术将解决你所有的问题。但可惜，事实并非如此。在下一节中，我们将回到上面演示的怪物错误场景，并介绍最后一种调试技术。

五、使用 TensorFlow 自定义训练循环进行调试

        在我们上面描述的场景中，经过几天的训练，模型的特定状态和特定的训练批处理样本的组合突然导致损失变成 NaN。

让我们评估如何使用上面的调试技术来调试此问题。

• 如果我们仔细跟踪用于所有随机操作的种子，并且没有不受控制的非确定性事件，理论上我们可以通过从头开始训练来重现错误......但这需要几天时间。
• 在本地环境或预先执行模式下复制可能需要数周时间。
• 我们可以从最近的检查点恢复，但是如果我们可以从完全相同的样本恢复并且所有伪随机生成器的状态完全相同，我们只能重现相同的模型状态和批处理样本。
• 添加 tf.prints 会有所帮助，但会带来巨大的开销
• 添加tf.debugging.enable_check_numerics对于查明失败的功能将非常有帮助。如果函数中存在明显的错误，这可能就足够了。但它不能让我们重现该错误。

        理想情况下，我们将能够在损失消失之前捕获输入和模型状态。然后，我们可以在受控（本地）环境中，以预先执行模式并使用调试器重现问题。

        问题是，我们不知道这个问题即将发生，直到它真正发生。当损失报告为 NaN 时，模型已使用 NaN 权重进行了更新，并且导致错误的批处理样本已经迭代。

        我想提出的解决方案是自定义训练循环，以便我们在每一步记录当前样本，并且仅在梯度有效时才更新模型权重。如果梯度无效，我们将停止训练并转储最后一个批次样本以及当前模型快照。这可以转移到本地环境，在那里加载模型，并在预先执行模式下输入捕获的数据样本，以便重现（和解决）错误。

        我们稍后将介绍代码，但首先，简要介绍使用自定义训练循环的利弊。

5.1 自定义训练循环与高级 API

        TensorFlow用户之间关于是编写自定义训练循环还是依赖高级API（如tf.keras.model.fit（））存在一个由来已久的争议。

        自定义培训循环的支持者预示着能够逐行控制培训的执行方式，以及发挥创造力的自由。高级 API 的支持者指出了它提供的许多便利，最值得注意的是内置的回调实用程序和分布式策略支持。还假定使用高级 API 可以确保您使用的是无错误且高度优化的训练循环实现。

        从 2.2 版本开始，TensorFlow 引入了覆盖 tf.keras.model 类的 train_step 和make_train_function例程的功能。这使人们能够引入一定程度的自定义，同时继续享受 model.fit（） 的便利。我们将演示如何以这样一种方式覆盖这些函数，使我们能够捕获有问题的示例输入和模型状态以进行本地调试。

5.2 自定义捕获循环

        在下面的代码块中，我们使用train_step和make_train_functions例程的自定义实现扩展了 tf.keras.models.Model 对象。为了全面了解实现，我建议您将其与 github 中例程的默认实现进行比较。您会注意到，我已经删除了与指标计算和策略支持相关的所有逻辑，以使代码更具可读性。需要注意的主要变化是：

• 在将梯度应用于模型权重之前，我们测试 NaN 的梯度。仅当 NaN 未出现时，渐变才会应用于权重。否则，将向训练循环发送遇到错误的信号。信号的一个例子是将损耗设置为预定值，例如零或NaN。
• 训练循环在每个步骤中存储数据特征和标签（x 和 y）。请注意，为了做到这一点，我们已将数据集遍历（下一个（迭代器）调用）移到了 @tf.function 范围之外。
• 该类有一个布尔“crash”标志，用于向 main 函数发出是否遇到错误的信号。

class CustomKerasModel(tf.keras.models.Model):def __init__(self, **kwargs):super(CustomKerasModel, self).__init__(**kwargs)# boolean flag that will signal to main function that # an error was encounteredself.crash = [e-mail] train_step(self, data):x, y = datawith tf.GradientTape() as tape:y_pred = self(x, training=True)  # Forward pass# Compute the loss value# (the loss function is configured in `compile()`)loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)# Compute gradientstrainable_vars = self.trainable_variablesgradients = tape.gradient(loss, trainable_vars)# concatenate the gradients into a single tensor for testingconcat_grads = tf.concat([tf.reshape(g,[-1]) for g in gradients],0)# In this example, we test for NaNs, # but we can include other testsif tf.reduce_any(tf.math.is_nan(concat_grads)):# if any of the gradients are NaN, send a signal to the  # outer loop and halt the training. We choose to signal# to the outer loop by setting the loss to 0.return {'loss': 0.}else:# Update weightsself.optimizer.apply_gradients(zip(gradients, trainable_vars))return {'loss': loss}def make_train_function(self):if self.train_function is not None:return self.train_functiondef train_function(iterator):data = next(iterator)# records the current sampleself.x, self.y = datares = self.train_step(data)if res['loss'] == 0.:self.crash = Trueraise Exception()return resself.train_function = train_functionreturn self.train_functionif __name__ == '__main__':# train_ds = # inputs = # outputs =# optimizer =# loss =  # epochs =# steps_per_epoch = model = CustomKerasModel(inputs=inputs, outputs=outputs)opt = tf.keras.optimizers.Adadelta(1.0)model.compile(loss=loss, optimizer=optimizer)try:model.fit(train_ds, epochs=epochs,                         steps_per_epoch=steps_per_epoch)except Exception as e:# check for signalif model.crash:model.save_weights('model_weights.ckpt')# pickle dump model.x and model.yfeatures_dict = {}for n, v in model.x.items():features_dict[n] = v.numpy()with open('features.pkl','wb') as f:pickle.dump(features_dict,f)labels_dict = {}for n, v in model.y.items():labels_dict[n] = v.numpy()with open('labels.pkl', 'wb') as f:pickle.dump(labels_dict, f)raise e

        请务必注意，此技术的训练运行时成本很小，它来自在预先执行模式（而不是图形模式）下从数据集读取数据。（没有免费的午餐。确切的成本将取决于模型的大小;模型越大，这种变化就越少。您应该在自己的模型上评估此技术的开销，然后决定是否以及如何使用它。

六、总结

        只要我们人类参与人工智能应用程序的开发，编程错误的普遍性就几乎是有保证的。在设计代码时考虑到可调试性，并获得解决错误的工具和技术，可以防止一些严重的折磨。

        最重要的是，不要气馁。