如何在两个长度不等的tf.dataset上进行成对迭代？

Question

我处理两个长度不等的数据集

我的目标是为datasetA中的每个元素从其他datasetB中获取一个元素。我尝试.take(1)（如图所示here）从datasetB获取单个元素，但重复调用.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的迭代器是很困难的

我进一步发现的是for old TensorFlow，它使用TF操作重新初始化迭代器，迭代器不再可用。 在this question中也提到了这一点，这很有帮助，但并没有让我找到TF2.+解决方案

我所寻找的是一种从datasetA和datasetB获取成对元素的优雅方法，其中datasetB在使用时（自动）重复

我不需要迭代组合数据集，除非较短的数据集通过重复“填充”到较长的数据集，然后我可以从datasetA中的a和datasetB中的B采样对（a，B）

TL；博士： 要在两个长度不等的数据集上进行成对迭代，请在使用较短的数据集时重新启动

Answer 1

要从两个数据集中获取所有可能的样本对，可以使用以下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这两个数据集