将tf.dataset写回TFRecord

Question

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

要做到这一点，一种方法是遍历整个dataset，并在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文件中的一行。

例如(从官方文档中保留) tfds使用下面的generate_examples来保存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。