问题排查

本指南为想要在 Cloud TPU 上运行自己的 TensorFlow 模型的用户提供了问题排查帮助。如需更广泛的 Cloud TPU 使用入门指南,请参阅快速入门MNIST 教程

概览

在 Cloud TPU 上运行 TensorFlow 模型的推荐策略是使用 TPUEstimator API。如果您目前正在使用 TensorFlow 的 Estimator API,则切换到 TPUEstimator 通常只需更改几行代码。将数据加载到 TPUEstimator 的推荐方法是使用 tf.data API。如需查看如何实现 TPUEstimatortf.data 的实际示例,请参阅将 TPUEstimator API 与 Cloud TPU 结合使用中的 MNIST 示例。另请参阅 MNIST Estimator 到 TPUEstimator

将您的模型转换为 TPUEstimator 后,请确保该模型支持 use_tpu=False 标志。将该标志设置为 false 会使 TensorFlow 回退到 Estimator API,不使用任何与 TPU 相关的代码。在 use_tpu=False 的情况下运行模型时遇到的任何问题都与 TPU 无关,并且不在本指南的讨论范围内。请改为参阅 TensorFlow 程序人员指南

理想情况下,如果模型已经可以使用 TPUEstimatoruse_tpu=False 成功运行,则在 TPU 上运行该模型只需设置 use_tpu=True 并将 master 指向一个 TPU 服务器网址(通常使用集群解析器来实现)。但是,由于 TensorFlow 模型可能非常复杂,并且 TPU 使用其专有的执行引擎,因此可能会遇到一些 TPU 特有的问题。这些问题分为以下几大类:

  1. 训练脚本完全无法连接到 TPU 服务器。

  2. 在尝试执行模型时 TPU 返回错误。

  3. TPU 内存无法容纳模型。

  4. 模型可以在 TPU 上运行,但达不到预期训练速度。

  5. 模型可以在 TPU 上运行,但 TPU 训练的模型的准确率低于 CPU/GPU 训练的模型的基准。

此外,本指南还包含有关 TPU 上可用的常规功能的常见问题解答

如需将特定类型的神经网络植入 TPU 方面的更加专业化的帮助,请参阅 Cloud TPU 教程

无法连接到 TPU 服务器

在 TPU 上运行模型时,您必须将远程 TPU 服务器网址传递给 RunConfig 中的 master 参数。从本质上讲,TensorFlow 创建一个与此服务器有关的远程 tf.Session。本部分提供了针对 TensorFlow 在连接到 TPU 服务器时停止响应或输出错误等情况的问题排查办法。请注意,大型模型的 TPU 图编译步骤可能需要很长时间,因此请先让脚本执行至少 5 分钟,然后再判断它是否已停止响应。

首先,验证问题在于服务器本身,还是在于 TensorFlow 训练流水线。为此,请使用您的 TPU 服务器网址运行 MNIST 教程并验证它可以正常运行。如果使用 MNIST 教程时仍然遇到连接问题,则表明问题在于 TPU 服务器。在这种情况下,请执行以下操作:

  1. 运行以下命令以列出可用的 TPU:

    (vm)$ gcloud compute tpus list
    

    您可能还需要设置 zoneproject,如 MNIST 教程中所示。这会显示如下输出:

    NAME       ZONE           ACCELERATOR_TYPE  NETWORK_ENDPOINT   NETWORK  RANGE          STATUS
    demo-tpu   us-central1-b  v2-8              10.240.1.2:8470    default  10.240.1.0  READY

  2. 验证您将正确的值传递给 --tpu(上例中的 demo-tpu),并且此 TPU 被列为 READY。此外,请确保按照如下方式设置您的 zoneproject

    (vm)$ gcloud config set project your-project-name
    
    (vm)$ gcloud config set compute/zone us-central1-b
    
  3. 如果此 TPU 未被列为 READY,或者您仍然无法建立连接,请使用 gcloud compute tpus stop $TPU_SERVER_NAME && gcloud compute tpus start $TPU_SERVER_NAME 手动重启服务器。在上面的示例中,$TPU_NAMEdemo-tpu。这可能需要几分钟。

  4. 重新运行上面的 ... tpus list 命令并等待 TPU 处于 READY 状态。这可能需要几分钟。

  5. 尝试再次运行 MNIST 教程。

  6. 如果您在运行 MNIST 教程时仍遇到问题,请使用获取支持中介绍的某种机制寻求帮助。

如果 MNIST 示例可正确运行但您的模型仍然停止响应,则问题可能出在您的训练流水线上。此时,请首先确保您的模型使用的是 TPUEstimator API,因为它不仅可以应对复杂的处理流水线,还支持使用 use_tpu 标志在 TPU 与非 TPU 执行之间轻松切换。请参阅 TPU 教程,了解几个有关如何使用 TPUEstimator 的示例。如果您的模型使用的是 TPUEstimator API,请在设置 use_tpu=False 后验证它能够正确运行。如果在设置 use_tpu=False 后模型无法正确运行,则问题与 TPU 无关。

调试常见错误

无法使用本地文件系统

错误消息

InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

详情

所有输入文件和模型目录都必须使用 Cloud Storage 存储分区路径 (gs://bucket-name/...),并且此存储分区必须可通过 TPU 服务器访问。请注意,所有数据处理和模型检查点都在 TPU 服务器上执行,而不是在本地机器上执行。如需了解如何正确配置 Cloud Storage 以便与 TPU 搭配使用,请参阅连接到 Cloud Storage 存储分区指南。

tf.data.Dataset.cache() 无法缓存到本地文件系统

错误消息

tensorflow.python.framework.errors_impl.UnimplementedError: File system scheme '[local]' not implemented (file: '[filename].lockfile')

详情

tf.data.Dataset 可以缓存。.cache() 调用有两个实现,分别位于:

  1. 内存中(如果没有传递任何参数)。

  2. 文件系统上(如果文件路径作为参数传递)。

在 Cloud TPU 上,(1) 适用(只要它适合可用的内存),但是当保存到本地文件系统时,(2) 不适用,并会导致上述错误。

以下代码段对这两种情况进行了说明:

(1)
 import tensorflow as tf

def main():
  print('Hello world!')
  ds = tf.data.Dataset.range(10)
  ds = ds.cache()

runs to completion.

(2)
 import tensorflow as tf

def main():
  print('Hello world!')
  ds = tf.data.Dataset.range(10)
  ds = ds.cache('/tmp/foo')

generates the error.

API 指南包含更多有关 tf.data.Dataset.cache() 的详细信息。

不支持的数据类型

错误消息

TypeError: DataType is not a supported TPU infeed type.

详情

目前,TPU 仅支持 tf.float32tf.int32tf.bfloat16tf.bool 数据类型。其他常见数据类型(例如 tf.uint8tf.stringtf.int64)必须在数据预处理期间(即在 TPUEstimatorinput_fn 中)转换为某种受支持的数据类型。请参阅 MNIST 教程查看另一示例。以 MNIST 中的以下代码段为例,它会将存储为 tf.uint8 字节序列的 image 张量转换为 tf.float32 张量:

image = tf.decode_raw(image, tf.uint8)
image = tf.cast(image, tf.float32)
image = tf.reshape(image, [784])

此代码段会将存储为 tf.int64label 张量转换为 tf.int32 张量:

label = tf.cast(label, tf.int32)

不支持动态形状

错误消息

ValueError: shape [Shape] must have a fixed size for dimension d that is known at graph construction time.

详情

如需在 TPU 上执行模型,TensorFlow 使用 XLA 框架编译该模型。虽然此编译步骤可显著提高训练速度和改进内存使用情况,但图中所有张量的形状(维度大小)都必须是静态的,也就是说,在进行图编译时它们的值必须是已知的。 如果有任何形状在编译时无法确定,则 TPU 编译将失败,并显示类似上面所示的错误。

返回动态形状的一个常用操作是 dataset.batch(batch_size),因为数据流中剩余的样本数可能小于批量大小。因此,在 TPU 上进行训练时,请使用 tf.contrib.data.batch_and_drop_remainder(batch_size)。这可能会丢弃文件中的最后几个样本,以确保每个批量都具有静态形状 (batch_size)。例如:

dataset = ...
dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(batch_size))

不可用的 TensorFlow 操作

错误消息

NotFoundError: No registered 'OpName' OpKernel for XLA_TPU_JIT devices compatible with node

详情

模型使用目前无法在 TPU 上执行的 TensorFlow 操作。

如需了解 TPU 上可用操作的列表,以及未来支持计划和解决方法建议,请参阅可用的 TensorFlow 操作指南。

内存不足错误消息

错误消息

ResourceExhaustedError: Ran out of memory in memory space hbm; used: YYY; limit: 7.48G.

详情

每个 Cloud TPU 都由 8 个 TPU 核心组成,每个核心都有 8GB 的 RAM(或 HBM,即高带宽内存)。这些内存用于存储权重(可变)张量,以及梯度计算所需的中间结果张量。如果模型太大而无法放入 TPU RAM 中,初始化将失败并显示上述错误消息。如需更多帮助,请参阅有关减少内存使用量的部分。

不使用 CrossShardOptimizer

错误消息

ValueError: CrossShardOptimizer must be used for model training on TPUs.

详情

使用 TensorFlow Python API 定义模型时,用户编写的绝大多数代码不需要针对 TPU 进行专门处理。最重要的一种例外情况是优化器,它必须按如下所示封装在 tf.contrib.tpu.CrossShardOptimizer() 中:

optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
if FLAGS.use_tpu:
  optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)
train_op=optimizer.minimize(loss, tf.train.get_global_step())

每个 Cloud TPU 都由 8 个 TPU 核心组成,这些核心是独立的处理单元。对于每个训练步(例如权重更新),每个 TPU 核心都针对独立的小批量数据运行前向传导和梯度计算,然后所有核心彼此交换梯度。尽管存在了解数据分片中介绍的一些注意事项,但在大多数情况下,这在数学上等同于计算一个大型批量的梯度。

CrossShardOptimizer 是负责执行此梯度交换的操作。默认情况下,CrossShardOptimizer 会计算所有核心上平均损失的梯度,但也可以通过传递 reduction=losses.Reduction.SUM 将其配置为计算损失总和

无法连接到 TPU 服务器

错误消息

An error was raised while a session was being created. This may be due to a preemption of a connected worker or parameter server. A new session is created.

详情

当 TensorFlow 无法连接到传递给 master 的 TPU 服务器网址时,系统会显示此错误。如需帮助,请参阅有关无法连接 TPU 服务器的部分。

训练过程中出现错误

如果模型无法在 TPU 上成功执行,则与此相关的所有错误都会在初始化期间被捕获。因此,模型在训练过程中失败的情况很少见。如果确实发生了这种情况,最可能的原因是数据预处理函数存在问题。例如,使用 Dataset API 时,通常需要调用 dataset = dataset.repeat(),否则在进行一次数据传递后,训练即会失败。像 tf.while_loop() 这样的动态执行操作同样仅会因为输入数据而失败。在极少数情况下,还可能发生虚假硬件或网络故障。

停止执行时出现问题

如果 TensorFlow 在 TPU 执行期间遇到错误,则脚本有时会停止响应,而不是退出 shell。如果发生这种情况,请按键盘上的 CTRL+\ 以触发 SIGQUIT,这会使 Python 立即退出。

同样,在 TPU 执行过程中按 CTRL+C 并不会立即关闭 TensorFlow,而是一直等到当前迭代循环结束才彻底退出。按 CTRL+\ 可使 Python 立即退出。

如果在以这种方式退出后,当重新连接到 TPU 服务器时遇到 DeadlineExceededError 等任何新错误,请使用 gcloud compute tpus stop $TPU_SERVER_NAME && gcloud compute tpus start $TPU_SERVER_NAME 命令(其中 $TPU_SERVER_NAME 取自 gcloud compute tpus list 命令的第一列)手动重置 TPU 服务器。

减少内存使用量

如果您在 TPU 上执行模型时遇到内存不足错误,则必须采取相应措施来减少该模型的内存使用量。本部分介绍了内存问题的几个根本原因,并提供解决这些问题的准则。

模型权重数量过多

内存问题的可能原因

每个 float32 模型权重需要 4 个字节。这些权重会在每个 TPU 核心上进行复制。因此,具有数亿个权重的模型可能因太大而无法放在 TPU 上。

如何减少内存使用量

  1. 某些优化器要求为每个权重提供额外的内存来存储更新统计信息。值得注意的是,AdamOptimizerAdadeltaOptimizer 都要求为每个权重额外提供 8 个字节。AdagradOptimizerMomentumOptimizer 都要求为每个权重额外提供 4 个字节。标准的 GradientDescentOptimizer 不需要任何额外的存储空间,但在最终模型准确率方面可能不如其他优化器。在训练 Transformer 模型时,实验性 AdafactorOptimizer 几乎不需要额外的内存,并且与基准 Adam 优化器的性能一样好。
  2. 如果大多数权重都是字词嵌入,则 WordPiece 之类的技术已被证实可以显著减少词量,同时提高执行多种任务的准确率。
  3. 即将发布的 TensorFlow 版本将提供对 16 位浮点权重和梯度的实验性支持,这会将内存需求减少一半。

过度张量填充

内存问题的可能原因

TPU 内存中的张量经过了填充,也就是说,TPU 将存储在内存中的张量的大小向上舍入,以便更高效地执行计算。此填充操作在硬件级别以透明方式进行,不会影响结果。但是,在某些情况下,填充可能会导致内存使用量和执行时间显著增加。

如何减少内存使用量

TPU 软件会尝试在内存中存放张量,以便最大限度地提高计算效率并减少填充。但是,此内存存放进程很复杂,为了获得最佳结果,该模型应遵循以下经验法则。为了最大限度地减少内存开销并提高计算效率,必须满足以下条件之一:

  • 总批量大小应为 64 的倍数(每个 TPU 核心 8 个),并且特征维度应为 128 的倍数。

    或者

  • 总批量大小应为 1024 的倍数(每个 TPU 核心 128 个),并且特征维度应为 8 的倍数。

使用 1024 作为批量大小以及 128 的倍数作为特征维度可以获得最佳效率,但这并非适用于所有模型。 为清楚起见,“特征维度”是指全连接层的隐藏大小或卷积中的输出通道的数量。并非所有层都符合此规则,尤其是网络的第一层和最后一层。这没什么问题,而且大多数模型预期都需要一定量的填充。

批量大小太大

内存问题的可能原因

在 CPU、GPU 或 TPU 上训练神经网络时,内存使用来自两个方面:

  1. 存储权重、权重梯度和特定于优化器的统计信息(例如动量)。内存使用量与模型中的权重数量(而不是批量大小)成正比。
  2. 存储来自于前向传导、计算反向传导所需的中间激活。内存使用量与批量大小、层大小和层数成正比。

综上所述,模型所需的内存在很大程度上取决于批量大小。

如何减少内存使用量

尝试缓慢减小批量大小,直到它能够存放在内存中为止,同时确保总批量大小为 64 的倍数(每个核心的批量大小应为 8 的倍数)。请注意,较大的批量大小在 TPU 上更为有效。通常,总批量大小为 1024(每个核心 128 个)是不错的起点。

模型过大

内存问题的可能原因

模型所需的内存在很大程度上取决于图中运算符的数量(即网络中的层数)。此存储要求与权重数量无关。例如,除了用于存储权重的任何内存之外,计算运算符(如 tf.nn.conv2d())的梯度也可能会增加内存使用量。

TPU 引擎会尝试策略性地重新计算某些运算符,以便使模型能够存放在内存中(称为再实质化 (rematerialization),类似于梯度检查点),但有时无法做到这一点。

如何减少内存使用量

如果即使批量较小(例如 64),模型也无法在 TPU 上运行,请尝试减少层数或层大小。 即将推出的 TensorFlow 版本将支持在 TPU 上“并行处理模型”,借助此支持,您可以在不同的 TPU 核心上运行模型的不同部分,从而实现在 Cloud TPU 上运行更大的模型。

提高训练速度

如果您的模型能够在 TPU 上成功运行,但训练速度低于预期,可以参考本部分概述的几种有助于提高训练速度的方法。

每个循环的迭代次数太少

性能问题描述

TPUConfigiterations_per_loop 参数可控制在单个“训练循环”中向 TPU 发送多少个数据批量。每个训练循环都需要在本地机器和 TPU 服务器之间进行大量通信,因此如果 iterations_per_loop 太小,则可能会大大减慢训练速度。

如何确定您的模型是否受到影响

如果系统非常频繁地(例如每 3 秒)显示日志记录消息 Enqueue next (X) batch(es) of data to infeed,那么您的训练可能在训练循环中产生了大量开销。

如何缓解

iterations_per_loop 设置为更大的值。在 MNIST 教程中,这由 --iterations 标志控制。只要 Enqueue next (X) batch(es) of data to infeed 消息的显示频率不超过每分钟几次,那么当前值应该就足够了。请注意,iterations_per_loop 可以设置为非常大的值,唯一的缺点是日志记录消息和检查点只能在循环结束时发生。

输入处理瓶颈

性能问题描述

当 TPU 在特定数据块上进行训练时,输入处理函数会在 CPU 上准备下一个数据块。因此,如果输入函数花费的时间少于模型函数,则输入处理的费用实际上为零。但是,如果输入函数花费的时间长于模型函数,则会产生瓶颈。

如何确定您的模型是否受到影响

按照 Cloud TPU 工具:输入流水线分析器中的说明在 TensorBoard 中查看流水线分析:

图片

输入流水线分析页面会显示一份清晰的摘要,其中显示了您的模型是否因输入处理而遭遇瓶颈。此页面还显示每个操作的执行时间,便于您找出有问题的操作。

如何缓解

使用 Dataset 加载数据时,有几种可能的缓解措施:

  1. 将您的数据以一组 tf.train.Example 结构的形式存储在 TFRecord 文件中,并使用 TFRecordDataset 加载它们。如需查看示例,请参阅 Dataset API 教程ResNet 教程
  2. 使用 dataset.cache() 和/或 dataset.prefetch() 来缓冲输入数据。这可以防止因文件访问中偶尔出现的速度缓慢问题而导致产生瓶颈。
  3. 指定 dataset.map() 函数的 num_parallel_calls 参数以启用多线程 map() 操作。
  4. 离线执行数据预处理,而不是在每次训练的每个周期执行预处理。此方法的优点是一劳永逸,缺点是离线执行数据预处理的费用高昂。

所有输入处理都是在位于 TPU 服务器上的 CPU 上执行,而不是在本地机器上执行,因此本地机器的速度不是一个影响因素。

非矩阵乘法操作过多

性能问题描述

Cloud TPU 能够以极快的速度执行矩阵乘法和卷积。大多数其他 TensorFlow 操作也可在 TPU 上高效地实现,但这些并不是 TPU 相对于其他硬件的主要优势。因此,模型应该以矩阵乘法或卷积为主,才能使 TPU 发挥最大作用。

如何确定您的模型是否受到影响

Cloud TPU 工具:操作配置文件指南介绍了如何为您的模型生成按操作类型细分的性能配置文件。通常,绝大多数现代神经网络架构都以矩阵乘法和卷积为主。

如何缓解

如果模型中缺少矩阵乘法主要是为了解决其他硬件上的训练速度问题,则建议您在 TPU 上重新对这些模型进行基准测试,以获得更高的速度。如果模型的基本属性就是缺少矩阵乘法,则 TPU 可能不是最佳的硬件选择。

过度张量填充

性能问题描述

TPU 会在内存中填充张量,以便高效地使用其计算单元。填充可能会增加内存和内存带宽的使用量。如需理解填充和解决张量填充问题方面的帮助,请参阅张量填充

批量大小太小

性能问题描述

一般来讲,使用较大的批量大小可提高 TPU 上的训练速度(以样本数/秒表示)。

如何确定您的模型是否受到影响

任何模型的批量大小应至少为 64(每个 TPU 核心 8 个),因为 TPU 总是会将张量填充到此大小。在 TPU 上训练时,理想的批量大小为 1024(每个 TPU 核心 128 字节),因为这样可以消除内存传输和填充导致的低效问题。

如何缓解

建议您使用可存放在内存中的最大批量大小,并且为 64 的倍数。实现此目的的最简单方法是从 1024 开始,如果该值会导致出现内存不足的错误,则尝试减小批量大小,直到模型成功运行为止。更改模型的批量大小可能需要调整其他超参数以实现同样的模型准确率(例如学习速率),但必须根据具体情况对其进行评估,不能一概而论。

层大小太小

性能问题描述

即使模型以矩阵乘法或卷积为主,如果输入张量很小,TPU 也可能无法全速运行。与其他硬件相比,当批量大小和层大小都很大时(例如维度 >= 512),TPU 运行的效率最高。

如何确定您的模型是否受到影响

一般来讲,小于 128 的层大小在 TPU 上的效率低下,因为 128 是 TPU 矩阵乘法单元的原生维度。对于全连接层,建议采用 512 作为隐藏大小下限,以实现高效率。注意,卷积层通常无需与全连接层一样大,即可实现相同的效率水平。例如,大小为 256 的 3×3 卷积可实现与大小为 2048 的全连接层相当的高效率,因为 3×3×256 = 2304。

如何缓解

如果您在模型中采用较小的层大小主要是为了训练速度,则建议您在 TPU 上使用更大的层重新对模型进行基准测试。例如,若将层的输出大小从 256 增加到 512,那么即使模型执行的计算量增加一倍,训练时间也只会增加 20%。

操作级模型分析

为了识别性能瓶颈,测量操作级执行时间和内存使用情况通常很有用。如需了解如何执行此操作的说明,
请参阅 Cloud TPU 工具:Trace Viewer 指南。

针对模型准确率降低进行调试

Cloud TPU 生态系统的目标之一是:目前在 CPU 或 GPU 上进行训练的任何模型在 TPU 上训练时都能达到非常接近的准确率,并且可能只需对批量大小和学习速率等超参数进行细微调整。但是,用户有时会发现在 TPU 上训练模型时准确率有所降低。由于神经网络训练的随机性,针对这些问题进行调试可能会让人很头痛。本部分提供一些指导,帮助您确定在将模型移植到 TPU 后模型准确率下降的根本原因。

了解数据分片(数据并行处理)

TensorFlow 的主要目标之一是每个操作应该产生几乎相同的结果,无论它是在 CPU、GPU 还是 TPU 上执行。但有一些例外情况,例如随机操作。一般来讲,如果您发现 TPU 与 CPU 上的非随机操作输出存在任何明显差异,请将其作为 bug 上报

但是,对于整个训练流水线,在 CPU/GPU 上与在 TPU 上进行训练存在一项重要的差异:当使用 TPUEstimator 并且 use_tpu=False 时,TensorFlow 会回退到其标准执行引擎。该引擎每步训练一个批量。但是,在实际 TPU 上进行训练时,TensorFlow 会执行数据分片(也称为“使用同步 SGD 的数据并行处理”)。这是因为每个 Cloud TPU 都由 8 个 TPU 核心组成,这些核心作为独立的处理单元运行。因此,对于训练中的每一步,每个 TPU 核心都会获得传递过来的一个数据批量、计算权重梯度、相互交换梯度,然后计算权重更新。默认情况下,系统会在所有核心上计算损失平均值,但也可以通过更改 CrossShardOptimizer 的参数来计算总损失。

如果计算模型的总损失可以通过独立的每样本损失的平均值(或总和)来计算,则此过程在数学上等同于针对单个大批量进行训练。最常见的非独立按样本操作是批量归一化,它单独针对每个核心批量运行。例如,如果总批量大小为 128,则每核心批量大小为 16,并且 8 个核心中的每个都对自己的 16 个样本执行批量归一化。我们已发现在某些情况下对小批量(例如小于 32 个)进行批量归一化会导致准确率降低。理想情况下,在 TPU 上训练时的总批量大小可能很大(例如 256 到 1024 个),因此这种大小的批量不会造成大问题。但是,如果这种批量大小太大而导致模型无法存放到内存中,则必须根据具体情况评估分片的影响。

由于分片带来的复杂性,针对模型准确率降低这一问题进行调试时,首先要运行确定性的单核 TPU 训练,并将其与在 CPU/GPU 上训练的模型进行比较。通常此操作可以快速完成,因为它不要求将模型训练至收敛。

确定性训练

针对模型准确率差异进行调试很困难,其中一个原因是 TensorFlow 在每次训练模型时都使用不同的权重初始化和数据重排。将训练过程修改为确定性的,可保证多次运行会产生几乎相同的模型。本部分演示如何确定性地运行 MNIST 教程:

  1. 通过在 CPU 上运行一步来生成初始检查点文件。这一步用于实现确定性权重初始化。这也可以通过播种变量初始化器来实现,但这样做难度更高。
# Run training for 1 step to create an initial checkpoint.
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/init_output \
  --random_seed=12345 \
  --iterations=1
  --train_steps=1
  1. 修改输入函数中的任何数据重排函数以使用随机种子。 MNIST 教程中已完成了这一步。这适用于输入数据处理操作,因为它们始终在 CPU 上运行。模型函数中的随机操作在 TPU 与 CPU 之间可能不具有确定性。例如:
# In the flag definitions
tf.flags.DEFINE_integer("batch_size", None, "Random seed for training")

# In the input_fn
if FLAGS.random_seed is not None:
dataset = dataset.shuffle(seed=FLAGS.random_seed)
  1. 在 CPU 上运行同一个模型两次,以验证训练具有确定性。请注意,训练必须运行合理的步数(例如 1000 步),但不需要运行到收敛,因为在 CPU 上进行此操作可能非常慢。

    由于 CPU 训练与单核 TPU 训练相当,因此请使用适合单个 TPU 核心的批量大小(通常是用完整批量大小除以 8)。TensorFlow 不保证各次运行之间的逐位确定性,但损失应该非常接近:
# Copy the initial weights
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_1
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_1
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_2
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_2

# Run 1
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

# Output 1
accuracy = 0.9910644, global_step = 1000, loss = 0.025323588

# Run 2
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

# Output 2
accuracy = 0.9910644, global_step = 1000, loss = 0.025323414

单核 TPU 训练

可以确定性地运行 MNIST 教程以后,下一步是在 TPU 上重现 CPU 训练的结果,使用单个 TPU 核心来确定问题是在于数据分片还是在于 TPU 执行引擎本身。

下面说明了如何针对 MNIST 教程执行单核训练和评估:

# Use the same weight initialization as the CPU
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/tpu_output

# Run training for 1000 steps
python mnist.py \
    --use_tpu=True \
    --master=$GRPC_SERVER \
    --train_file=${STORAGE_BUCKET}/data/train.tfrecords \
    --model_dir=${STORAGE_BUCKET}/tpu_output \
    --random_seed=12345 \
    --batch_size=128 \
    --train_steps=1000 \
    --eval_steps=10

  accuracy = 0.9910644, global_step = 1000, loss = 0.02514153

损失与 CPU 训练的模型并不完全一致,但应该很接近。 如果您的模型的损失不接近,则可能表明您的 TPU 执行引擎中出现了 bug。在提交 bug 报告之前,请核实以下事项:

  1. 您为 TPUConfig 传递的是 num_shards=1

  2. 您的模型函数中没有任何随机操作,并且输入函数中的所有随机操作都已使用了正确的种子。

  3. 您为 CPU 和 TPU 训练使用的是相同的初始检查点文件。

调试多核 TPU 训练

如果您的模型在 CPU 与单核 TPU 上实现了相同的损失,那么问题可能属于以下某一种:

(a) 降级是由于采用不同初始化训练神经模型时的自然随机方差所导致。

(b) 降级是由于与 TPU 上的数据分片相关的问题所导致。

为了确定问题是否属于 (a),使用相同的权重初始化在 CPU/GPU 和多核 TPU 上重新训练完整模型(如上所述)可能会有帮助。

如果您确信准确率的降低具备统计显著性,那么最可能与数据分片相关的问题有:

  1. 如果您的模型以每个样本误差之和的形式计算损失,您可能需要将 reduction=losses.Reduction.SUM 传递给 CrossShardOptimizer。默认情况下,CrossShardOptimizer 会计算损失的平均值,而不是总和。
  2. 如果您的模型使用批量归一化,则总批量大小小于 256(即每个核心少于 32 个)可能会降低准确率。
  3. 如果您的模型包含逐批量损失函数,那么这将受到分片的影响。此类损失函数通常非常专业化。例如,2017 年 Karras 等人在训练生成对抗网络时使用了批量判别器。