Tensorflow：如何从 CPU tf.data.Dataset (from_generator) 预取 GPU 上的数据答案

【问题标题】：Tensorflow: How to prefetch data on the GPU from CPU tf.data.Dataset (from_generator)Tensorflow：如何从 CPU tf.data.Dataset (from_generator) 预取 GPU 上的数据
【发布时间】：2020-09-30 17:04:01
【问题描述】：

我正在努力解决以下问题。我正在使用 from_generator 方法创建一个 tf.data.Dataset。我在 CPU 上执行这些操作，因为我不想让我的 GPU 内存过载。

数据集由元组组成，其中包含一个固定长度的 tf.bool 一维掩码 (tf.Tensor) 和一个可变大小的 tf.float 二维矩阵 (tf.Tensor)。损失函数使用以下装饰器进行装饰，因此我不会认为可变大小是问题。

@tf.function(experimental_relax_shapes=True)

理想情况下，数据集保存在 CPU 上，然后预取到 GPU 上。

        def gen():
            for i, j in zip(mask_list, wmat_list):
                yield i, j

        dataset = tf.data.Dataset.from_generator(gen, output_types=(tf.bool, tf.float32))

目前主要的训练循环依赖 tf.identity 将数据移动到 gpu，效率低下。如下面 Tensorboard 的屏幕截图所示。大约 70% 的时间用于加载数据并将其移动到 GPU。

                for b, (mask, wmat) in enumerate(dataset):
                    with tf.GradientTape() as tape:

                        mask = tf.identity(mask)
                        wmat = tf.identity(wmat)

                        mean_error, loss = self.model.loss(mask, wmat)
                        epoch_loss += loss.numpy()
                        epoch_mean_error += mean_error.numpy()

我已经尝试过“prefetch_to_device”功能。但是，它没有将数据移动到 GPU 上。通过打印验证，例如训练循环中的 mask.device。

        gpu_transform = tf.data.experimental.prefetch_to_device('/gpu')
        dataset.apply(gpu_transform)

对我来说，它类似于这个错误：https://github.com/tensorflow/tensorflow/issues/30929。但是，它被标记为已解决并且已经超过一年了。

使用官方 Docker 镜像运行 TF 2.3。

【问题讨论】：

标签： tensorflow tensorflow2.0

【解决方案1】：

我找到了自己问题的解决方案。

问题在于数据集中的元组确实不包含 tf.Tensors，而是 numpy 数组。因此，管道可能受到 py_func() 功能的限制。

下面的屏幕截图显示管道没有在 CPU 上阻塞。但是仍然有相当大的 MemCpy。 prefetch_to_device() 仍然没有做任何事情。这可能是由于一个已知问题应该在 TF2.4 中修复

https://github.com/tensorflow/tensorflow/issues/35563

~~（未经证实的）建议的解决方法对我也不起作用。~~（见编辑）

with tf.device("/gpu:0"):
    ds = ds.prefetch(1)

编辑：

我已进一步调查此问题并提交了错误报告。现在看来，建议的解决方法确实做了一些事情，但不确定它是否完全及时预取。 https://github.com/tensorflow/tensorflow/issues/43905

【讨论】：