使用 tf.estimator.train_and_evaluate 时如何有效地洗牌大型 tf.data.Dataset？

Question

该tf.estimator.train_and_evaluate文件明确指出，输入数据集必须正确洗牌的训练，看看所有的例子：

过拟合：为了避免过拟合，建议设置训练 input_fn 对训练数据进行适当的 shuffle。还建议在执行评估之前对模型进行更长时间的训练，比如多个时期，因为每次训练的输入管道都从头开始。这对于本地培训和评估尤为重要。

在我的应用程序中，我想从tf.data.Dataset具有任意评估频率和shuffle()缓冲区大小的完整样本中统一采样示例。否则，训练最多只能看到第一个：

(steps_per_second * eval_delay * batch_size) + buffer_size

元素，有效地丢弃其余元素。有没有一种有效的方法来解决这个问题，而无需在系统内存中加载完整的数据集？

我考虑过根据缓冲区大小对数据集进行分片，但如果评估不经常发生，它将在同一个分片上迭代多次（一个repeat()关闭管道）。理想情况下，我想在对数据集进行完整迭代后移动到另一个分片，这可能吗？

感谢您的任何指点！

Answer 1

可以通过此数据集转换来实现数据集的随机分片：

def random_shard(shard_size, dataset_size):
  num_shards = -(-dataset_size // shard_size)  # Ceil division.
  offsets = np.linspace(
      0, dataset_size, num=num_shards, endpoint=False, dtype=np.int64)

  def _random_shard(dataset):
    sharded_dataset = tf.data.Dataset.from_tensor_slices(offsets)
    sharded_dataset = sharded_dataset.shuffle(num_shards)
    sharded_dataset = sharded_dataset.flat_map(
        lambda offset: dataset.skip(offset).take(shard_size))
    return sharded_dataset

  return _random_shard

这需要提前知道数据集的总大小。但是，如果您实现基于文件的分片方法，您还可以对完整数据集进行一次迭代，因此这不是一个主要问题。

关于效率，请注意，skip(offset)实际上会迭代示例，因此如果很大，offset则预计会出现延迟。offset仔细的预取应该对此有所帮助。