将 tf.dataset 写回 TFRecord

Question

创建 tf.data.Dataset 后，我​​想将其写入 TFRecords。

一种方法是遍历整个数据集并在 serializeToString 之后写入 TFRecords。但这并不是最有效的方法。

有没有更简单的方法来做到这一点？ TF2.0 中是否有可用的 API？

Answer 1

您可以使用TensorFlow Datasets (tfds)：这个库不仅是一个现成可用的tf.data.Dataset 对象的集合，而且还是一个将原始数据转换为TFRecords 的工具链。

按照official guide 可以直接添加新数据集。总之，你只需要实现_info和_generate_examples这两个方法。

特别是，_generate_examples 是 tfds 用来在 TFRecords 中创建行的方法。 _generate_examples 产生的每个元素都是字典；每个字典都是 TFRecord 文件中的一行。

例如（从官方文档中保留）下面的generate_examples是tfds用来保存TFRecords的，每一个都有记录“image_description”、“image”、“label”。

def _generate_examples(self, images_dir_path, labels):
  # Read the input data out of the source files
  for image_file in tf.io.gfile.listdir(images_dir_path):
    ...
  with tf.io.gfile.GFile(labels) as f:
    ...

  # And yield examples as feature dictionaries
  for image_id, description, label in data:
    yield image_id, {
        "image_description": description,
        "image": "%s/%s.jpeg" % (images_dir_path, image_id),
        "label": label,
    }

在您的情况下，您可以只使用已有的 tf.data.Dataset 对象，并循环遍历它（在 generate_examples 方法中），并产生 TFRecord 的行。

通过这种方式，tfds 会为您处理序列化，您会在 ~/tensorflow_datasets 文件夹中找到为您的数据集创建的 TFRecord。