在生成分片 tfrecord 时创建循环分片答案

【问题标题】：Create round-robin sharding while generating sharded tfrecords在生成分片 tfrecord 时创建循环分片
【发布时间】：2019-08-23 19:15:32
【问题描述】：

我是 tensorflow 的新手，我正在处理 tensorflow 1.14 中的图像分割问题。当我尝试生成一个大的 tfrecord 文件时，我有一个巨大的数据集并且生成 tfrecord 的速度非常慢。所以，我想创建 'n' 个 tfrecords 碎片。我找不到在网上做的方法。假设我有 600 张图像和 600 个面具。我想以循环方式生成 6 个 tfrecord 分片，每个分片有 100 个图像和 100 个掩码。我想要的高级/伪代码如下-

sharded_tf_record_writer:
create n TFRecordWriter
----> for each_item in n TFRecordWriter
      -----> write_example in round-robin fashion

我确实在网上搜索过，但找不到相关答案。我不想使用 apache 光束进行分片。我感谢实现这一目标的任何想法/帮助/指导。

【问题讨论】：

标签： tensorflow computer-vision image-segmentation sharding tfrecord

【解决方案1】：

我在 tensorflow 数据集和用户的问题之一中问过同样的问题 - Conchylicultor 对此做出了回应 -

写入由 _TFRecordWriter 完成。 Tfds 将自动计算所需的分片数量并在分片之间分配示例，但是每个分片都是按顺序写入的。您无法控制分片的数量，它也会自动计算。

但是，示例分布在分片之间的事实并不能加快编写速度，因为示例不是并行预处理的。如果你想要并行性，那么你必须使用 Apache Beam，它甚至可以扩展到巨大的数据集

张量流/数据集问题的链接是 - https://github.com/tensorflow/datasets/issues/676

这可能会有所帮助。

【讨论】：

【解决方案2】：

由于您正在使用 tensorflow 中的对象检测，因此官方 Tensorflow models 存储库中有一些不错的代码可以满足您的需求。请注意，此代码适用于 Tensorflow2（不确定它是否适用于 TF1）

参见 example 从 coco 注释中编写分片 tfrecords。这个想法是您在退出堆栈中打开一个 TFRecordWriter 列表（使用contextlib2.ExitStack()），当每个线程完成写入时，它将自动关闭 TFRecords。

实用函数open_sharded_output_tfrecords 函数创建这个 TFRecordWriter 列表

import contextlib2
import tensorflow as tf
with contextlib2.ExitStack() as tf_record_close_stack, tf.gfile.GFile(
    annotations_file, 'r'
) as fid:
    output_tfrecords = tf_record_creation_util.open_sharded_output_tfrecords(
        tf_record_close_stack, output_path, num_shards
    )

接下来，您可以使用 ProcessPoolExecutor 以循环方式并行将 tfrecords 写入每个分片（本例中为 4 个工作人员）

from concurrent.futures.process import ProcesPoolExecutor
with ProcessPoolExecutor(4) as executor:
    for idx, image in enumerate(images):
        futures = []
        future = executor.submit(
            _write_tf_record,
            image,
            idx,
            num_shards,
            output_tfrecords,
        )
        futures.append(future)
    for future in futures:
        future.result()

_write_tf_record 可能看起来像这样：

def _write_tf_record(image, idx, num_shards, output_tfrecords)
    tf_example = create_tf_example(image)
    shard_idx = idx % num_shards
    output_tfrecords[shard_idx].write(tf_example.SerializeToString())

只要确保你的分片比多进程工作者多，否则同一个 writer 可能会被两个不同的进程访问。

【讨论】：