TFRecord 和 tf.Example

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

TFRecord 和 tf.Example#

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

TFRecord 格式是一种用于存储二进制记录序列的简单格式。

协议缓冲区是一个跨平台、跨语言的库，用于高效地序列化结构化数据。

协议消息由 .proto 文件定义，这通常是了解消息类型最简单的方法。

tf.train.Example 消息（或 protobuf）是一种灵活的消息类型，表示 {"string": value} 映射。它可以与 TensorFlow 结合使用，并在 TFX 等更高级 API 中使用。

此笔记本将演示如何创建、解析和使用 tf.Example 消息，以及如何在 .tfrecord 文件之间对 tf.Example 消息进行序列化、写入和读取。

注：这些结构虽然有用，但并不是强制的。您无需转换现有代码即可使用 TFRecord，除非您正在使用 tf.data 且读取数据仍是训练的瓶颈。有关数据集性能的提示，请参阅使用 tf.data API 提升性能。

注：通常，您应当将数据分片到多个文件，以便可以并行化 I/O（在单个主机内或跨多个主机）。根据经验法则，文件数量至少应达到读取数据的主机数量的 10 倍。同时，每个文件都应当足够大（至少 10+MB，理想情况下为 100MB+），以便您从 I/O 预提取中受益。例如，假设您有 X GB 数据，并且您计划在最多 N 个主机上进行训练。理想情况下，您应当将数据分片到 ~10*N 个文件，只要 ~X/(10*N) 为 10+ MB（理想情况下为 100+ MB）。如果小于该值，则可能需要创建更少的分片来权衡并行性优势和 I/O 预提取优势。

设置#

import tensorflow as tf

import numpy as np
import IPython.display as display

`tf.train.Example`#

`tf.Example` 的数据类型#

从根本上讲，tf.Example 是 {"string": tf.train.Feature} 映射。

tf.train.Feature 消息类型可以接受以下三种类型（请参阅 .proto 文件）。大多数其他通用类型也可以强制转换成下面的其中一种：

tf.train.BytesList（可强制转换自以下类型）

string
byte

tf.train.FloatList（可强制转换自以下类型）

float (float32)
double (float64)

tf.train.Int64List（可强制转换自以下类型）

bool
enum
int32
uint32
int64
uint64

为了将标准 TensorFlow 类型转换为兼容 tf.Example 的 tf.train.Feature，可以使用下面的快捷函数。请注意，每个函数会接受标量输入值并返回包含上述三种 list 类型之一的 tf.train.Feature：

# The following functions can be used to convert a value to a type compatible
# with tf.train.Example.

def _bytes_feature(value):
  """Returns a bytes_list from a string / byte."""
  if isinstance(value, type(tf.constant(0))):
    value = value.numpy() # BytesList won't unpack a string from an EagerTensor.
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _float_feature(value):
  """Returns a float_list from a float / double."""
  return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))

def _int64_feature(value):
  """Returns an int64_list from a bool / enum / int / uint."""
  return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

注：为了简单起见，本示例仅使用标量输入。要处理非标量特征，最简单的方法是使用 tf.io.serialize_tensor 将张量转换为二进制字符串。在 TensorFlow 中，字符串是标量。使用 tf.io.parse_tensor 可将二进制字符串转换回张量。

下面是有关这些函数如何工作的一些示例。请注意不同的输入类型和标准化的输出类型。如果函数的输入类型与上述可强制转换的类型均不匹配，则该函数将引发异常（例如，_int64_feature(1.0) 将出错，因为 1.0 是浮点数，应该用于 _float_feature 函数）：

print(_bytes_feature(b'test_string'))
print(_bytes_feature(u'test_bytes'.encode('utf-8')))

print(_float_feature(np.exp(1)))

print(_int64_feature(True))
print(_int64_feature(1))

可以使用 .SerializeToString 方法将所有协议消息序列化为二进制字符串：

feature = _float_feature(np.exp(1))

feature.SerializeToString()

创建 `tf.Example` 消息#

假设您要根据现有数据创建 tf.Example 消息。在实践中，数据集可能来自任何地方，但是从单个观测值创建 tf.Example 消息的过程相同：

在每个观测结果中，需要使用上述其中一种函数，将每个值转换为包含三种兼容类型之一的 tf.train.Feature。
创建一个从特征名称字符串到第 1 步中生成的编码特征值的映射（字典）。
将第 2 步中生成的映射转换为 Features 消息。

在此笔记本中，您将使用 NumPy 创建一个数据集。

此数据集将具有 4 个特征：

具有相等 False 或 True 概率的布尔特征
从 [0, 5] 均匀随机选择的整数特征
通过将整数特征作为索引从字符串表生成的字符串特征
来自标准正态分布的浮点特征

请思考一个样本，其中包含来自上述每个分布的 10,000 个独立且分布相同的观测值：

# The number of observations in the dataset.
n_observations = int(1e4)

# Boolean feature, encoded as False or True.
feature0 = np.random.choice([False, True], n_observations)

# Integer feature, random from 0 to 4.
feature1 = np.random.randint(0, 5, n_observations)

# String feature.
strings = np.array([b'cat', b'dog', b'chicken', b'horse', b'goat'])
feature2 = strings[feature1]

# Float feature, from a standard normal distribution.
feature3 = np.random.randn(n_observations)

您可以使用 _bytes_feature、_float_feature 或 _int64_feature 将下面的每个特征强制转换为兼容 tf.Example 的类型。然后，可以通过下面的已编码特征创建 tf.Example 消息：

def serialize_example(feature0, feature1, feature2, feature3):
  """
  Creates a tf.train.Example message ready to be written to a file.
  """
  # Create a dictionary mapping the feature name to the tf.train.Example-compatible
  # data type.
  feature = {
      'feature0': _int64_feature(feature0),
      'feature1': _int64_feature(feature1),
      'feature2': _bytes_feature(feature2),
      'feature3': _float_feature(feature3),
  }

  # Create a Features message using tf.train.Example.

  example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
  return example_proto.SerializeToString()

例如，假设您从数据集中获得了一个观测值 [False, 4, bytes('goat'), 0.9876]。您可以使用 create_message() 创建和打印此观测值的 tf.Example 消息。如上所述，每个观测值将被写为一条 Features 消息。请注意，tf.Example 消息只是 Features 消息外围的包装器：

# This is an example observation from the dataset.

example_observation = []

serialized_example = serialize_example(False, 4, b'goat', 0.9876)
serialized_example

要解码消息，请使用 tf.train.Example.FromString 方法。

example_proto = tf.train.Example.FromString(serialized_example)
example_proto

TFRecords 格式详细信息#

TFRecord 文件包含一系列记录。该文件只能按顺序读取。

每条记录包含一个字节字符串（用于数据有效负载），外加数据长度，以及用于完整性检查的 CRC32C（使用 Castagnoli 多项式的 32 位 CRC）哈希值。

每条记录会存储为以下格式：

uint64 length
uint32 masked_crc32_of_length
byte   data[length]
uint32 masked_crc32_of_data

将记录连接起来以生成文件。此处对 CRC 进行了说明，且 CRC 的掩码为：

masked_crc = ((crc &gt;&gt; 15) | (crc &lt;&lt; 17)) + 0xa282ead8ul

注：没有要求在 TFRecord 文件中使用 tf.train.Example。tf.train.Example 只是一种将字典序列化为字节串的方法。任何可以在 TensorFlow 中解码的字节串都可以存储在 TFRecord 文件中。示例包括：文本行、JSON（使用 tf.io.decode_json_example）、编码图像数据或序列化 tf.Tensors（使用 tf.io.serialize_tensor/tf.io.parse_tensor）。请参阅 tf.io 模块了解更多选项。

使用 `tf.data` 的 TFRecord 文件#

tf.data 模块还提供用于在 TensorFlow 中读取和写入数据的工具。

写入 TFRecord 文件#

要将数据放入数据集中，最简单的方式是使用 from_tensor_slices 方法。

若应用于数组，将返回标量数据集：

tf.data.Dataset.from_tensor_slices(feature1)

若应用于数组的元组，将返回元组的数据集：

features_dataset = tf.data.Dataset.from_tensor_slices((feature0, feature1, feature2, feature3))
features_dataset

# Use `take(1)` to only pull one example from the dataset.
for f0,f1,f2,f3 in features_dataset.take(1):
  print(f0)
  print(f1)
  print(f2)
  print(f3)

使用 tf.data.Dataset.map 方法可将函数应用于 Dataset 的每个元素。

映射函数必须在 TensorFlow 计算图模式下进行运算（它必须在 tf.Tensors 上运算并返回）。可以使用 tf.py_function 包装非张量函数（如 serialize_example）以使其兼容。

使用 tf.py_function 需要指定形状和类型信息，否则它将不可用：

def tf_serialize_example(f0,f1,f2,f3):
  tf_string = tf.py_function(
    serialize_example,
    (f0, f1, f2, f3),  # Pass these args to the above function.
    tf.string)      # The return type is `tf.string`.
  return tf.reshape(tf_string, ()) # The result is a scalar.

tf_serialize_example(f0, f1, f2, f3)

将此函数应用于数据集中的每个元素：

serialized_features_dataset = features_dataset.map(tf_serialize_example)
serialized_features_dataset

def generator():
  for features in features_dataset:
    yield serialize_example(*features)

serialized_features_dataset = tf.data.Dataset.from_generator(
    generator, output_types=tf.string, output_shapes=())

serialized_features_dataset

并将它们写入 TFRecord 文件：

filename = 'test.tfrecord'
writer = tf.data.experimental.TFRecordWriter(filename)
writer.write(serialized_features_dataset)

读取 TFRecord 文件#

您还可以使用 tf.data.TFRecordDataset 类来读取 TFRecord 文件。

有关通过 tf.data 使用 TFRecord 文件的详细信息，请参见此处。

使用 TFRecordDataset 对于标准化输入数据和优化性能十分有用。

filenames = [filename]
raw_dataset = tf.data.TFRecordDataset(filenames)
raw_dataset

此时，数据集包含序列化的 tf.train.Example 消息。迭代时，它会将其作为标量字符串张量返回。

使用 .take 方法仅显示前 10 条记录。

注：在 tf.data.Dataset 上进行迭代仅在启用了 Eager Execution 时有效。

for raw_record in raw_dataset.take(10):
  print(repr(raw_record))

可以使用以下函数对这些张量进行解析。请注意，这里的 feature_description 是必需的，因为数据集使用计算图执行，并且需要以下描述来构建它们的形状和类型签名：

# Create a description of the features.
feature_description = {
    'feature0': tf.io.FixedLenFeature([], tf.int64, default_value=0),
    'feature1': tf.io.FixedLenFeature([], tf.int64, default_value=0),
    'feature2': tf.io.FixedLenFeature([], tf.string, default_value=''),
    'feature3': tf.io.FixedLenFeature([], tf.float32, default_value=0.0),
}

def _parse_function(example_proto):
  # Parse the input `tf.train.Example` proto using the dictionary above.
  return tf.io.parse_single_example(example_proto, feature_description)

或者，使用 tf.parse example 一次解析整个批次。使用 tf.data.Dataset.map 方法将此函数应用于数据集中的每一项：

parsed_dataset = raw_dataset.map(_parse_function)
parsed_dataset

使用 Eager Execution 在数据集中显示观测值。此数据集中有 10,000 个观测值，但只会显示前 10 个。数据会作为特征字典进行显示。每一项都是一个 tf.Tensor，此张量的 numpy 元素会显示特征的值：

for parsed_record in parsed_dataset.take(10):
  print(repr(parsed_record))

在这里，tf.parse_example 函数会将 tf.Example 字段解压缩为标准张量。

Python 中的 TFRecord 文件#

tf.io 模块还包含用于读取和写入 TFRecord 文件的纯 Python 函数。

写入 TFRecord 文件#

接下来，将 10,000 个观测值写入文件 test.tfrecord。每个观测值都将转换为一条 tf.Example 消息，然后被写入文件。随后，您可以验证是否已创建 test.tfrecord 文件：

# Write the `tf.train.Example` observations to the file.
with tf.io.TFRecordWriter(filename) as writer:
  for i in range(n_observations):
    example = serialize_example(feature0[i], feature1[i], feature2[i], feature3[i])
    writer.write(example)

!du -sh {filename}

读取 TFRecord 文件#

您可以使用 tf.train.Example.ParseFromString 轻松解析以下序列化张量：

filenames = [filename]
raw_dataset = tf.data.TFRecordDataset(filenames)
raw_dataset

for raw_record in raw_dataset.take(1):
  example = tf.train.Example()
  example.ParseFromString(raw_record.numpy())
  print(example)

这将返回 tf.train.Example proto，它难以按照原样使用，但基本上可以表示：

Dict[str,
     Union[List[float],
           List[int],
           List[str]]]

以下代码可以手动将 Example 转换成 NumPy 数组的字典，无需使用 TensorFlow Ops。有关详情，请参阅 PROTO 文件。

result = {}
# example.features.feature is the dictionary
for key, feature in example.features.feature.items():
  # The values are the Feature objects which contain a `kind` which contains:
  # one of three fields: bytes_list, float_list, int64_list

  kind = feature.WhichOneof('kind')
  result[key] = np.array(getattr(feature, kind).value)

result

演练：读取和写入图像数据#

下面是关于如何使用 TFRecord 读取和写入图像数据的端到端示例。您将使用图像作为输入数据，将数据写入 TFRecord 文件，然后将文件读取回来并显示图像。

如果您想在同一个输入数据集上使用多个模型，这种做法会很有用。您可以不以原始格式存储图像，而是将图像预处理为 TFRecord 格式，然后将其用于所有后续的处理和建模中。

首先，让我们下载雪中的猫的图像，以及施工中的纽约威廉斯堡大桥的照片。

提取图像#

cat_in_snow  = tf.keras.utils.get_file(
    '320px-Felis_catus-cat_on_snow.jpg',
    'https://storage.googleapis.com/download.tensorflow.org/example_images/320px-Felis_catus-cat_on_snow.jpg')

williamsburg_bridge = tf.keras.utils.get_file(
    '194px-New_East_River_Bridge_from_Brooklyn_det.4a09796u.jpg',
    'https://storage.googleapis.com/download.tensorflow.org/example_images/194px-New_East_River_Bridge_from_Brooklyn_det.4a09796u.jpg')

display.display(display.Image(filename=cat_in_snow))
display.display(display.HTML('Image cc-by: <a "href=https://commons.wikimedia.org/wiki/File:Felis_catus-cat_on_snow.jpg">Von.grzanka</a>'))

display.display(display.Image(filename=williamsburg_bridge))
display.display(display.HTML('<a "href=https://commons.wikimedia.org/wiki/File:New_East_River_Bridge_from_Brooklyn_det.4a09796u.jpg">From Wikimedia</a>'))

写入 TFRecord 文件#

和以前一样，将特征编码为与 tf.Example 兼容的类型。这将存储原始图像字符串特征，以及高度、宽度、深度和任意 label 特征。后者会在您写入文件以区分猫和桥的图像时使用。将 0 用于猫的图像，将 1 用于桥的图像：

image_labels = {
    cat_in_snow : 0,
    williamsburg_bridge : 1,
}

# This is an example, just using the cat image.
image_string = open(cat_in_snow, 'rb').read()

label = image_labels[cat_in_snow]

# Create a dictionary with features that may be relevant.
def image_example(image_string, label):
  image_shape = tf.io.decode_jpeg(image_string).shape

  feature = {
      'height': _int64_feature(image_shape[0]),
      'width': _int64_feature(image_shape[1]),
      'depth': _int64_feature(image_shape[2]),
      'label': _int64_feature(label),
      'image_raw': _bytes_feature(image_string),
  }

  return tf.train.Example(features=tf.train.Features(feature=feature))

for line in str(image_example(image_string, label)).split('\n')[:15]:
  print(line)
print('...')

请注意，所有特征现在都存储在 tf.Example 消息中。接下来，函数化上面的代码，并将示例消息写入名为 images.tfrecords 的文件：

# Write the raw image files to `images.tfrecords`.
# First, process the two images into `tf.train.Example` messages.
# Then, write to a `.tfrecords` file.
record_file = 'images.tfrecords'
with tf.io.TFRecordWriter(record_file) as writer:
  for filename, label in image_labels.items():
    image_string = open(filename, 'rb').read()
    tf_example = image_example(image_string, label)
    writer.write(tf_example.SerializeToString())

!du -sh {record_file}

读取 TFRecord 文件#

现在，您有文件 images.tfrecords，并可以迭代其中的记录以将您写入的内容读取回来。因为在此示例中您只需重新生成图像，所以您只需要原始图像字符串这一个特征。使用上面描述的 getter 方法（即 example.features.feature['image_raw'].bytes_list.value[0]）提取该特征。您还可以使用标签来确定哪个记录是猫，哪个记录是桥：

raw_image_dataset = tf.data.TFRecordDataset('images.tfrecords')

# Create a dictionary describing the features.
image_feature_description = {
    'height': tf.io.FixedLenFeature([], tf.int64),
    'width': tf.io.FixedLenFeature([], tf.int64),
    'depth': tf.io.FixedLenFeature([], tf.int64),
    'label': tf.io.FixedLenFeature([], tf.int64),
    'image_raw': tf.io.FixedLenFeature([], tf.string),
}

def _parse_image_function(example_proto):
  # Parse the input tf.train.Example proto using the dictionary above.
  return tf.io.parse_single_example(example_proto, image_feature_description)

parsed_image_dataset = raw_image_dataset.map(_parse_image_function)
parsed_image_dataset

从 TFRecord 文件中恢复图像：

for image_features in parsed_image_dataset:
  image_raw = image_features['image_raw'].numpy()
  display.display(display.Image(data=image_raw))

TFRecord 和 tf.Example

目录

TFRecord 和 tf.Example#

设置#

tf.train.Example#

tf.Example 的数据类型#

创建 tf.Example 消息#

TFRecords 格式详细信息#

使用 tf.data 的 TFRecord 文件#

写入 TFRecord 文件#

读取 TFRecord 文件#

Python 中的 TFRecord 文件#

写入 TFRecord 文件#

读取 TFRecord 文件#

演练：读取和写入图像数据#

提取图像#

写入 TFRecord 文件#

读取 TFRecord 文件#

`tf.train.Example`#

`tf.Example` 的数据类型#

创建 `tf.Example` 消息#

使用 `tf.data` 的 TFRecord 文件#