时间戳

首页 > 解决方案 > 如何对两个不等长的 tf.datasets 进行成对迭代？

python - 如何对两个不等长的 tf.datasets 进行成对迭代？

问题描述

我使用两个长度不等的数据集。

我的目标是为 datasetA 中的每个元素获取另一个 datasetB 中的元素。我尝试.take(1)（如图所示）从数据集B 中获取单个元素，但重复调用.take(1)不会提前数据集的内部计数，即它始终返回相同的元素；但我每次都想获得一个新元素。

我可以使用 迭代一个数据集for element in datasetA:，然后将其中的第二个数据集作为elementB = iterB.get_next(). iterB这会在消耗时引发错误。

这是我正在使用的完整玩具代码：

datasetA = tf.data.Dataset.from_tensor_slices([1,  2,  3,  4,  5, 6])
datasetB = tf.data.Dataset.from_tensor_slices([11, 22, 33, 44])

iterB = iter(datasetB)
epochs = 5

for epoch in range(epochs):
  print(f"Epoch {epoch}")
  for element in datasetA:
    print(element)
    elementB = iterB.get_next()
    print(elementB)

然后我继续：

for epoch in range(epochs):
  print(f"Epoch {epoch}")
  for element in datasetA:
    print(element)
    elementB = iterB.get_next_as_optional()
    if not elementB.has_value():
      iterB = iter(datasetB)
      elementB = iterB.get_next_as_optional()

    print(elementB.get_value())

哪个有效，但是重新初始化datasetB.

我进一步发现的是旧的 TensorFlow，它使用 TF 操作来重新初始化不再可用的迭代器。这个问题中也提到了这一点，这很有帮助，但并没有引导我找到 TF2.+ 解决方案。

我正在寻找的是一种从datasetAand获取成对元素的优雅方法datasetB，其中 datasetB 在消耗时（自动）重复。

我不需要迭代一个组合数据集，除非较短的数据集通过重复被“填充”到较长的数据集，然后我可以对（A，B）与数据集A中的A和数据集B中的B进行采样。

TL;DR：想要对两个长度不等的数据集进行成对迭代，在使用时重新启动较短的数据集。

标签： pythontensorflowiterator

解决方案

要从两个数据集中获取所有可能的样本对，可以使用以下方法generator：

    # assuming that dataset_A and dataset_B are defined globally
    def generator():
        for sample_A in dataset_A:
            for sample_B in dataset_B:
                yield (sample_A, sample_B)

要仅获取数据集中相同位置（相同索引）的样本对，有一种标准zip方法：

    dataset = tf.data.Dataset.zip(dataset_A, dataset_B)

当其中一个数据集耗尽时，此类生成器将停止。

如果目标是从 中获取所有样本的成对样本dataset_A，但dataset_B更小，则可以无限重复第二个数据集，

     dataset_B = dataset_B.repeat()

然后zip是两个数据集。

推荐阅读