python - 如何在 tf.data 中启用渴望模式

Question

对于一个项目，我使用 tf.data.Dataset 来编写输入管道。输入是图像 RGB 标签是图像中用于生成热图的对象的 2D 坐标列表

这是重现问题的MWE。

 def encode_images(image, label):
        """

        Parameters
        ----------
        image
        label

        Returns
        -------

        """
        # load image
        # here the normal code
        # img_contents = tf.io.read_file(image)
        # # decode the image
        # img = tf.image.decode_jpeg(img_contents, channels=3)
        # img = tf.image.resize(img, (256, 256))
        # img = tf.cast(img, tf.float32)

        # this is just for testing
        image = tf.random.uniform(
            (256, 256, 3), minval=0, maxval=255, dtype=tf.dtypes.float32, seed=None, name=None
        )
        return image, label

    def generate_heatmap(image, label):
        """

        Parameters
        ----------
        image
        label

        Returns
        -------

        """

        start = 0.5
        sigma=3
        img_shape = (image.shape[0] , image.shape[1] )
        density_map = np.zeros(img_shape, dtype=np.float32)
        for center_x, center_y in label[0]:
            for v_y in range(img_shape[0]):
                for v_x in range(img_shape[1]):
                    x = start + v_x
                    y = start + v_y
                    d2 = (x - center_x) * (x - center_x) + (y - center_y) * (y - center_y)
                    exp = d2 / (2.0 * sigma**2)
                    if exp > math.log(100):
                        continue
                    density_map[v_y, v_x] = math.exp(-exp)
        return density_map


    X = ["img1.png", "img2.png", "img3.png", "img4.png", "img5.png"]
    y = [[[2, 3], [100, 120], [100, 120]],
         [[2, 3], [100, 120], [100, 120], [2, 1]],
         [[2, 3], [100, 120], [100, 120], [10, 10], [11, 12]],
         [[2, 3], [100, 120], [100, 120], [10, 10], [11, 12], [10, 2]],
         [[2, 3], [100, 120], [100, 120]]
         ]
    dataset = tf.data.Dataset.from_tensor_slices((X, tf.ragged.constant(y)))
    dataset = dataset.map(encode_images, num_parallel_calls=8)
    dataset = dataset.map(generate_heatmap, num_parallel_calls=8)
    dataset = dataset.batch(1, drop_remainder=False)

问题是在generate_heatmap()函数中，我使用 numpy 数组通过索引修改元素，这在 tensorflow 中是不可能的。我尝试迭代标签张量，这在 tensorflow 中是不可能的。另一件事是急切模式似乎没有启用 tf.data.Dataset！有什么建议可以解决这个问题！我认为在 pytorch 中这样的代码可以快速完成而不会受苦:)！

标签： pythonnumpytensorflowtf.data.dataset