TPU 速度数据流水线：tf.data.Dataset 和 TFRecords

1. 概览

TPU 的速度非常快。训练数据流必须跟上训练速度。在本实验中，您将学习如何使用 tf.data.Dataset API 从 GCS 加载数据，以馈送给 TPU。

本实验是“Keras on TPU”系列实验的第 1 部分。您可以按以下顺序执行这些步骤，也可以单独执行。

[本实验] TPU 速度数据流水线：tf.data.Dataset 和 TFRecords
您的第一个 Keras 模型，采用迁移学习
使用 Keras 和 TPU 的卷积神经网络
现代卷积神经网络、squeezenet、Xception 以及 Keras 和 TPU

学习内容

使用 tf.data.Dataset API 加载训练数据
使用 TFRecord 格式从 GCS 高效加载训练数据

反馈

如果您在此 Codelab 中发现任何问题，请告诉我们。您可以通过 GitHub 问题 [反馈链接] 提供反馈。

2. Google Colaboratory 快速入门

本实验使用 Google Colaboratory，无需您进行任何设置。Colaboratory 是一个在线笔记本平台，可用于教育目的。它提供免费的 CPU、GPU 和 TPU 训练。

您可以打开此示例笔记本并运行几个单元格，以便熟悉 Colaboratory。

Welcome to Colab.ipynb

选择 TPU 后端

在 Colab 菜单中，依次选择运行时 > 更改运行时类型，然后选择 TPU。在此代码实验中，您将使用强大的 TPU（张量处理单元）来支持硬件加速训练。首次执行时，系统会自动连接到运行时，您也可以使用右上角的“连接”按钮。

笔记本执行

点击某个单元格，然后使用 Shift-Enter 逐个执行单元格。您还可以通过依次选择运行时 > 全部运行来运行整个笔记本

隐藏单元格

部分单元格将仅显示其标题。这是 Colab 特有的笔记本功能。您可以双击它们来查看其中的代码，但通常不会很有趣。通常是支持或可视化函数。您仍需运行这些单元格，以便定义其中的函数。

Authentication

只要您使用已获授权的账号进行身份验证，Colab 就可以访问您的私有 Google Cloud Storage 存储分区。上述代码段将触发身份验证流程。

3. [信息] 什么是张量处理单元 (TPU)？

简而言之

在 Keras 中训练 TPU 模型的代码（如果 TPU 不可用，则回退到 GPU 或 CPU）：

try: # detect TPUs
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
    strategy = tf.distribute.TPUStrategy(tpu)
except ValueError: # detect GPUs
    strategy = tf.distribute.MirroredStrategy() # for CPU/GPU or multi-GPU machines

# use TPUStrategy scope to define model
with strategy.scope():
  model = tf.keras.Sequential( ... )
  model.compile( ... )

# train model normally on a tf.data.Dataset
model.fit(training_dataset, epochs=EPOCHS, steps_per_epoch=...)

我们今天将使用 TPU 以交互式速度（每次训练运行只需几分钟）构建和优化花卉分类器。

为什么选择 TPU？

现代 GPU 围绕可编程的“核心”构建，这种非常灵活的架构使它们能够处理各种任务，例如 3D 渲染、深度学习、物理模拟等。另一方面，TPU 将经典矢量处理器与专用矩阵乘法单元配对，擅长处理以大型矩阵乘法为主的任何任务，例如神经网络。

图示：密集神经网络层作为矩阵乘法，一次通过神经网络处理一批八张图片。请运行一行 x 列乘法，以验证它是否确实在对图像的所有像素值进行加权求和。卷积层也可以表示为矩阵乘法，不过会稍微复杂一些（ 相关说明请参阅第 1 部分）。

硬件

MXU 和 VPU

TPU v2 核心由一个用于运行矩阵乘法的矩阵乘法单元 (MXU) 和一个用于执行所有其他任务（例如激活、softmax 等）的向量处理单元 (VPU) 组成。VPU 可处理 float32 和 int32 计算。另一方面，MXU 以混合精度 16-32 位浮点格式运行。

混合精度浮点和 bfloat16

MXU 使用 bfloat16 输入和 float32 输出计算矩阵乘法。中间累积以 float32 精度执行。

神经网络训练通常能够抵抗因浮点精度降低而引入的噪声。在某些情况下，噪声甚至有助于优化器收敛。16 位浮点精度传统上用于加速计算，但 float16 和 float32 格式的范围差异很大。将精度从 float32 降低到 float16 通常会导致上溢和下溢。虽然有解决方案，但通常需要额外的工作才能使 float16 正常运行。

因此，Google 在 TPU 中引入了 bfloat16 格式。bfloat16 是一种截断的 float32，具有与 float32 完全相同的指数位和范围。此外，TPU 以混合精度计算矩阵乘法，输入为 bfloat16，输出为 float32，这意味着，通常无需更改代码即可受益于精度降低带来的性能提升。

脉动阵列

MXU 在硬件中实现矩阵乘法，采用所谓的“脉动阵列”架构，其中数据元素流经硬件计算单元阵列。（在医学上，“收缩”是指心脏收缩和血流（量），这里是指数据流。）

矩阵乘法的基本要素是一个矩阵中的一行与另一个矩阵中的一列之间的点积（请参阅本部分顶部的图示）。对于矩阵乘法 Y=X*W，结果的一个元素将为：

Y[2,0] = X[2,0]*W[0,0] + X[2,1]*W[1,0] + X[2,2]*W[2,0] + ... + X[2,n]*W[n,0]

在 GPU 上，人们会将此点积编程到 GPU“核心”中，然后并行地在尽可能多的“核心”上执行该点积，以尝试一次性计算出结果矩阵的每个值。如果生成的矩阵大小为 128x128，则需要 128x128=16K 个“核心”，这通常是不可能的。最大的 GPU 拥有大约 4000 个核心。另一方面，TPU 仅使用 MXU 中计算单元所需的最低限度的硬件：只有 bfloat16 x bfloat16 => float32 个乘法累加器，没有其他任何东西。这些矩阵非常小，TPU 可以在 128x128 MXU 中实现 16K 个这样的矩阵，并一次性处理此矩阵乘法。

图示：MXU 脉动阵列。计算元素是乘法累加器。一个矩阵的值会加载到数组中（红点）。其他矩阵的值会流经数组（灰色圆点）。竖线将值向上传播。横线用于传播部分和。用户可以自行验证，当数据流经数组时，您会从右侧获得矩阵乘法的结果。

此外，在 MXU 中计算点积时，中间和仅在相邻的计算单元之间流动。它们无需存储和检索到内存或寄存器文件。最终结果是，在计算矩阵乘法时，TPU 脉动阵列架构在密度和能耗方面具有显著优势，在速度方面也具有不可忽略的优势。

Cloud TPU

当您在 Google Cloud Platform 上请求一个“Cloud TPU v2”时，您会获得一个具有 PCI 连接的 TPU 板的虚拟机 (VM)。TPU 板有四个双核 TPU 芯片。每个 TPU 核心都包含一个 VPU（向量处理单元）和一个 128x128 MXU（矩阵乘法单元）。然后，此“Cloud TPU”通常通过网络连接到请求它的虚拟机。因此，完整的情况如下所示：

图示：您的虚拟机，其中包含一个网络连接的“Cloud TPU”加速器。“Cloud TPU”本身由一个虚拟机组成，该虚拟机具有一个 PCI 连接的 TPU 板，该板上包含四个双核 TPU 芯片。

TPU Pod

在 Google 的数据中心内，TPU 连接到高性能计算 (HPC) 互连，这使得它们看起来像一个非常大的加速器。Google 将它们称为 Pod，它们最多可包含 512 个 TPU v2 核心或 2048 个 TPU v3 核心。

图示：TPU v3 pod。通过 HPC 互连连接的 TPU 板和机架。

在训练期间，TPU 核心之间会使用全缩减算法交换梯度（此处对全缩减进行了很好的说明）。正在训练的模型可以通过以大批次大小进行训练来利用硬件。

图示：在 Google TPU 的二维环面网状网 HPC 网络上，使用 all-reduce 算法在训练期间同步梯度。

软件

大批次大小训练

TPU 的理想批次大小为每个 TPU 核心 128 个数据项，但硬件从每个 TPU 核心 8 个数据项开始就能显示出良好的利用率。请注意，一个 Cloud TPU 有 8 个核心。

在此 Codelab 中，我们将使用 Keras API。在 Keras 中，您指定的批次是整个 TPU 的全局批次大小。系统会自动将批次拆分为 8 个，并在 TPU 的 8 个核心上运行。

如需了解其他性能提示，请参阅 TPU 性能指南。对于非常大的批次大小，某些模型可能需要特别注意，请参阅 LARSOptimizer 了解详情。

深入了解：XLA

TensorFlow 程序定义计算图。TPU 不会直接运行 Python 代码，而是运行由 TensorFlow 程序定义的计算图。在底层，一个名为 XLA（加速线性代数编译器）的编译器会将计算节点的 TensorFlow 图转换为 TPU 机器代码。此编译器还会对您的代码和内存布局执行许多高级优化。当工作发送到 TPU 时，系统会自动进行编译。您不必在构建链中明确添加 XLA。

图示：为了在 TPU 上运行，Tensorflow 程序定义的计算图首先会被转换为 XLA（加速线性代数编译器）表示法，然后由 XLA 编译为 TPU 机器码。

在 Keras 中使用 TPU

从 TensorFlow 2.1 开始，Keras API 支持 TPU。Keras 支持可在 TPU 和 TPU Pod 上运行。以下示例适用于 TPU、GPU 和 CPU：

try: # detect TPUs
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
    strategy = tf.distribute.TPUStrategy(tpu)
except ValueError: # detect GPUs
    strategy = tf.distribute.MirroredStrategy() # for CPU/GPU or multi-GPU machines

# use TPUStrategy scope to define model
with strategy.scope():
  model = tf.keras.Sequential( ... )
  model.compile( ... )

# train model normally on a tf.data.Dataset
model.fit(training_dataset, epochs=EPOCHS, steps_per_epoch=...)

在此代码段中：

TPUClusterResolver().connect() 在网络上查找 TPU。在大多数 Google Cloud 系统（AI Platform 作业、Colaboratory、Kubeflow、通过“ctpu up”实用程序创建的深度学习虚拟机）上，该函数无需参数即可运行。这些系统通过 TPU_NAME 环境变量知道其 TPU 的位置。如果您手动创建 TPU，请在您要使用的虚拟机上设置 TPU_NAME 环境变量，或使用显式参数调用 TPUClusterResolver：TPUClusterResolver(tp_uname, zone, project)
TPUStrategy 是实现分布式和“全归约”梯度同步算法的部分。
策略通过范围应用。必须在 strategy 范围() 内定义模型。
tpu_model.fit 函数需要一个 tf.data.Dataset 对象作为 TPU 训练的输入。

常见的 TPU 移植任务

虽然有很多方法可以在 TensorFlow 模型中加载数据，但对于 TPU，必须使用 tf.data.Dataset API。
TPU 的运行速度非常快，因此在 TPU 上运行时，提取数据往往会成为瓶颈。您可以使用一些工具来检测数据瓶颈，并参考 TPU 性能指南中的其他性能提示。
int8 或 int16 数字会被视为 int32。TPU 没有以小于 32 位运行的整数硬件。
不支持某些 TensorFlow 操作。点击此处可查看列表。好消息是，此限制仅适用于训练代码，即模型的前向和后向传递。您仍然可以在数据输入流水线中使用所有 TensorFlow 操作，因为这些操作将在 CPU 上执行。
TPU 不支持 tf.py_func。

4. 加载数据

我们将使用一个鲜花图片数据集。目标是学习将它们分为 5 种花卉类型。数据加载是使用 tf.data.Dataset API 执行的。首先，让我们了解一下该 API。

具备实践性

请打开以下笔记本，执行单元格（按 Shift-Enter 键），然后按照标有“需要完成的工作”的说明操作。

Fun with tf.data.Dataset (playground).ipynb

其他信息

“flowers”数据集简介

数据集整理在 5 个文件夹中。每个文件夹都包含一种花。这些文件夹的名称分别为 sunflowers、daisy、dandelion、tulips 和 roses。数据托管在 Google Cloud Storage 的公共存储分区中。摘录：

gs://flowers-public/sunflowers/5139971615_434ff8ed8b_n.jpg
gs://flowers-public/daisy/8094774544_35465c1c64.jpg
gs://flowers-public/sunflowers/9309473873_9d62b9082e.jpg
gs://flowers-public/dandelion/19551343954_83bb52f310_m.jpg
gs://flowers-public/dandelion/14199664556_188b37e51e.jpg
gs://flowers-public/tulips/4290566894_c7f061583d_m.jpg
gs://flowers-public/roses/3065719996_c16ecd5551.jpg
gs://flowers-public/dandelion/8168031302_6e36f39d87.jpg
gs://flowers-public/sunflowers/9564240106_0577e919da_n.jpg
gs://flowers-public/daisy/14167543177_cd36b54ac6_n.jpg

为什么使用 tf.data.Dataset？

Keras 和 TensorFlow 在其所有训练和评估函数中都接受数据集。在数据集中加载数据后，该 API 会提供所有对神经网络训练数据有用的常见功能：

dataset = ... # load something (see below)
dataset = dataset.shuffle(1000) # shuffle the dataset with a buffer of 1000
dataset = dataset.cache() # cache the dataset in RAM or on disk
dataset = dataset.repeat() # repeat the dataset indefinitely
dataset = dataset.batch(128) # batch data elements together in batches of 128
AUTOTUNE = tf.data.AUTOTUNE
dataset = dataset.prefetch(AUTOTUNE) # prefetch next batch(es) while training

您可以在这篇文章中找到性能提示和数据集最佳实践。点击此处查看参考文档。

tf.data.Dataset 基础知识

数据通常以多个文件（此处为图片）的形式呈现。您可以通过调用以下函数来创建文件名数据集：

filenames_dataset = tf.data.Dataset.list_files('gs://flowers-public/*/*.jpg')
# The parameter is a "glob" pattern that supports the * and ? wildcards.

然后，您将一个函数“映射”到每个文件名，该函数通常会将文件加载并解码为内存中的实际数据：

def decode_jpeg(filename):
  bits = tf.io.read_file(filename)
  image = tf.io.decode_jpeg(bits)
  return image

image_dataset = filenames_dataset.map(decode_jpeg)
# this is now a dataset of decoded images (uint8 RGB format)

如需迭代处理数据集，请执行以下操作：

for data in my_dataset:
  print(data)

元组数据集

在监督学习中，训练数据集通常由训练数据和正确答案对组成。为了实现这一点，解码函数可以返回元组。然后，您将获得一个元组数据集，当您对其进行迭代时，系统会返回元组。返回的值是 TensorFlow 张量，可供模型直接使用。您可以对它们调用 .numpy() 以查看原始值：

def decode_jpeg_and_label(filename):
  bits = tf.read_file(filename)
  image = tf.io.decode_jpeg(bits)
  label = ... # extract flower name from folder name
  return image, label

image_dataset = filenames_dataset.map(decode_jpeg_and_label)
# this is now a dataset of (image, label) pairs 

for image, label in dataset:
  print(image.numpy().shape, label.numpy())

结论：逐个加载图片速度很慢！

随着您对此数据集进行迭代，您会发现每秒可以加载大约 1-2 张图片。这太慢了！我们将用于训练的硬件加速器可以维持比此速率高出许多倍的速率。请前往下一部分，了解我们将如何实现此目标。

解决方案

以下是解决方案笔记本。如果您遇到困难，可以使用它。

Fun with tf.data.Dataset (solution).ipynb

所学内容

🤔 tf.data.Dataset.list_files
🤔 tf.data.Dataset.map
🤔 元组数据集
😀 遍历数据集

请花点时间在脑海中过一遍此核对清单。

5. 快速加载数据

我们将在本实验中使用的张量处理单元 (TPU) 硬件加速器速度非常快。挑战往往在于如何快速地向它们提供数据，以确保它们始终处于忙碌状态。Google Cloud Storage (GCS) 能够维持非常高的吞吐量，但与所有云存储系统一样，建立连接需要一些网络往返时间。因此，将数据存储为数千个单独的文件并不理想。我们将它们分批放入数量较少的文件中，并利用 tf.data.Dataset 的强大功能并行读取多个文件。

读穿

以下笔记本中包含用于加载图片文件、将它们调整为共同大小，然后将它们存储到 16 个 TFRecord 文件中的代码。请快速浏览一下。由于本 Codelab 的其余部分将提供格式正确的 TFRecord 数据，因此无需执行此操作。

Flower pictures to TFRecords.ipynb

实现最佳 GCS 吞吐量的理想数据布局

TFRecord 文件格式

TensorFlow 存储数据的首选文件格式是基于 protobuf 的 TFRecord 格式。其他序列化格式也适用，但您可以直接通过以下代码从 TFRecord 文件加载数据集：

filenames = tf.io.gfile.glob(FILENAME_PATTERN)
dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.map(...) # do the TFRecord decoding here - see below

为获得最佳性能，建议使用以下更复杂的代码同时从多个 TFRecord 文件中读取数据。此代码将并行读取 N 个文件，并忽略数据顺序以提高读取速度。

AUTOTUNE = tf.data.AUTOTUNE
ignore_order = tf.data.Options()
ignore_order.experimental_deterministic = False

filenames = tf.io.gfile.glob(FILENAME_PATTERN)
dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTOTUNE)
dataset = dataset.with_options(ignore_order)
dataset = dataset.map(...) # do the TFRecord decoding here - see below

TFRecord 备忘单

TFRecord 中可以存储三种类型的数据：字节字符串（字节列表）、64 位整数和 32 位浮点数。它们始终以列表形式存储，单个数据元素将是一个大小为 1 的列表。您可以使用以下辅助函数将数据存储到 TFRecord 中。

写入字节串

# warning, the input is a list of byte strings, which are themselves lists of bytes
def _bytestring_feature(list_of_bytestrings):
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=list_of_bytestrings))

写入整数

def _int_feature(list_of_ints): # int64
  return tf.train.Feature(int64_list=tf.train.Int64List(value=list_of_ints))

写入浮点数

def _float_feature(list_of_floats): # float32
  return tf.train.Feature(float_list=tf.train.FloatList(value=list_of_floats))

使用上述辅助函数写入 TFRecord

# input data in my_img_bytes, my_class, my_height, my_width, my_floats
with tf.python_io.TFRecordWriter(filename) as out_file:
  feature = {
    "image": _bytestring_feature([my_img_bytes]), # one image in the list
    "class": _int_feature([my_class]),            # one class in the list
    "size": _int_feature([my_height, my_width]),  # fixed length (2) list of ints
    "float_data": _float_feature(my_floats)       # variable length  list of floats
  }
  tf_record = tf.train.Example(features=tf.train.Features(feature=feature))
  out_file.write(tf_record.SerializeToString())

如需从 TFRecord 读取数据，您必须先声明所存储记录的布局。在声明中，您可以将任何已命名字段作为固定长度列表或可变长度列表进行访问：

从 TFRecord 读取

def read_tfrecord(data):
  features = {
    # tf.string = byte string (not text string)
    "image": tf.io.FixedLenFeature([], tf.string), # shape [] means scalar, here, a single byte string
    "class": tf.io.FixedLenFeature([], tf.int64),  # shape [] means scalar, i.e. a single item
    "size": tf.io.FixedLenFeature([2], tf.int64),  # two integers
    "float_data": tf.io.VarLenFeature(tf.float32)  # a variable number of floats
  }

  # decode the TFRecord
  tf_record = tf.io.parse_single_example(data, features)

  # FixedLenFeature fields are now ready to use
  sz = tf_record['size']

  # Typical code for decoding compressed images
  image = tf.io.decode_jpeg(tf_record['image'], channels=3)

  # VarLenFeature fields require additional sparse.to_dense decoding
  float_data = tf.sparse.to_dense(tf_record['float_data'])

  return image, sz, float_data

# decoding a tf.data.TFRecordDataset
dataset = dataset.map(read_tfrecord)
# now a dataset of triplets (image, sz, float_data)

有用的代码段：

读取单个数据元素

tf.io.FixedLenFeature([], tf.string)   # for one byte string
tf.io.FixedLenFeature([], tf.int64)    # for one int
tf.io.FixedLenFeature([], tf.float32)  # for one float

读取固定大小的元素列表

tf.io.FixedLenFeature([N], tf.string)   # list of N byte strings
tf.io.FixedLenFeature([N], tf.int64)    # list of N ints
tf.io.FixedLenFeature([N], tf.float32)  # list of N floats

读取数量不定的数据项

tf.io.VarLenFeature(tf.string)   # list of byte strings
tf.io.VarLenFeature(tf.int64)    # list of ints
tf.io.VarLenFeature(tf.float32)  # list of floats

VarLenFeature 会返回稀疏向量，并且在解码 TFRecord 后需要执行额外的步骤：

dense_data = tf.sparse.to_dense(tf_record['my_var_len_feature'])

还可以在 TFRecord 中添加可选字段。如果您在读取字段时指定了默认值，那么如果该字段缺失，系统会返回默认值，而不是错误。

tf.io.FixedLenFeature([], tf.int64, default_value=0) # this field is optional

所学内容

🤔 对数据文件进行分片，以便从 GCS 快速访问
😓 如何编写 TFRecord。（您已经忘记了语法？没关系，您可以将此页面添加为书签，以便日后参考）
🤔 使用 TFRecordDataset 从 TFRecord 加载数据集

请花点时间在脑海中过一遍此核对清单。

6. 恭喜！

您现在可以向 TPU 提供数据了。请继续进行下一个实验

[本实验] TPU 速度数据流水线：tf.data.Dataset 和 TFRecords
您的第一个 Keras 模型，采用迁移学习
使用 Keras 和 TPU 的卷积神经网络
现代卷积神经网络、squeezenet、Xception 以及 Keras 和 TPU

TPU 实践

Cloud AI Platform 提供 TPU 和 GPU：

最后，我们非常欢迎您提供反馈。如果您发现此实验有任何问题，或者认为此实验应进行改进，请告诉我们。您可以通过 GitHub 问题 [反馈链接] 提供反馈。

Martin Görner ID small.jpg
作者：Martin Görner
Twitter：@martin_gorner

TPU 速度数据流水线：tf.data.Dataset 和 TFRecords

1. 概览

学习内容

反馈

2. Google Colaboratory 快速入门

选择 TPU 后端

笔记本执行

目录

隐藏单元格

Authentication

3. [信息] 什么是张量处理单元 (TPU)？

简而言之

为什么选择 TPU？

硬件

MXU 和 VPU

混合精度浮点和 bfloat16

脉动阵列

Cloud TPU

TPU Pod

软件

大批次大小训练

深入了解：XLA

在 Keras 中使用 TPU

常见的 TPU 移植任务

4. 加载数据

具备实践性

其他信息

“flowers”数据集简介

为什么使用 tf.data.Dataset？

tf.data.Dataset 基础知识

元组数据集

结论：逐个加载图片速度很慢！

解决方案

所学内容

5. 快速加载数据

读穿

实现最佳 GCS 吞吐量的理想数据布局

TFRecord 文件格式

TFRecord 备忘单

所学内容

6. 恭喜！

TPU 实践