将tf.dataset中的每个样本映射为id

Question

出于测试目的，我想在tf.dataset中的每个样本上附加一个ID。简单地向上计数就足够了。

我的数据集的类型为FlatMapDataset fwiw。

for entry in img_ds:
        print(entry.shape)

（（128，128，3）（128，128，3）（128，128，3）（128，128，3）...

我尝试的是具有一个映射函数，该函数在其中定义一个计数器并向上计数：

@staticmethod
    def map_to_id(img):
        try:
            ExperimentalPipeline.map_to_id.id_counter += 1
        except AttributeError:
            ExperimentalPipeline.map_to_id.id_counter = 0
        return img, ExperimentalPipeline.map_to_id.id_counter

然后使用tf.data中的Dataset.map将ID附加到每个样本：

img_ds = img_ds.map(ExperimentalPipeline.map_to_id)

不幸的是，这不起作用，每个样本的ID为零：

for i, id in img_ds:
        print(f"{i.shape}, {id}")

（（128，128，3），0（128，128，3），0（128，128，3），0（128，128，3），0...

我还注意到我的map_to_id函数仅被调用一次。

@staticmethod
def map_to_id(img):
    print("enter map_to_id")
    try:
        ExperimentalPipeline.map_to_id.id_counter += 1
    except AttributeError:
        print("caught exception")
        ExperimentalPipeline.map_to_id.id_counter = np.random.randint(1000)
    return img, ExperimentalPipeline.map_to_id.id_counter

输入map_to_id捕获的异常（128，128，3），889（128，128，3），889（128，128，3），889（128，128，3），889

我想我不理解Dataset.map应该如何工作。我虽然会获取被调用的数据集中的每个样本，然后以该样本为参数调用提供的函数。有人可以帮我解决这个问题吗？

Answer 1

TensorFlow将运行一次map函数，以将该函数编译为TensorFlow操作。然后，这些操作（而不是原始的python函数）将应用于数据集的每个元素。如果要为每个元素运行原始的python函数，则可以改用py_function。

在这种情况下，您想附加元素ID，可以使用Dataset.enumerate实现您的目标：

img_ds = img_ds.enumerate()

Answer 2

好吧，阅读更多tensorflow文档后，我发现了这一点：

请注意，与定义map_func的上下文无关（渴望与图表），tf.data会跟踪该函数并将其作为图形。要在函数内部使用Python代码，您有两个选择：

1）依靠AutoGraph将Python代码转换为等效图计算。这种方法的缺点是AutoGraph可以转换一些但不是全部的Python代码。

2）使用tf.py_function，它允许您编写任意Python代码但通常会导致性能不及1）

因此map_to_id函数实际上仅被跟踪一次。由于选项1）似乎不起作用，所以我只选择选项2）。我只需要一些单元测试的ID，因此性能应该不是问题。

解决方案如下：

img_ds = img_ds.map(
    lambda img: tf.py_function(
        func=ExperimentalPipeline.map_to_id, inp=[img], Tout=(tf.float32, tf.int32)
    )
)