调试 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

2023-11-07 19:59:22.275296: 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'