Python 使用tf.map\u fn将多个图像作为张量读取

我正在使用Tensorflow的Dataset API阅读各种图像（数据和标签）。由于数据集队列位于CPU上，因此四处复制数据的成本很高。然而，我似乎无法找到一种方法来避免这种情况

问题：我能否以统一顺序（例如h、w、c）高效加载各种图像

假设我想在一个单通道图像中读取，我可以按如下方式执行：

image = tf.image.decode_png(tf.read_file(file_name), channels=1)  # h,w,c

或对于多通道RGB：

image = tf.image.decode_png(tf.read_file(file_name), channels=3)  # h,w,c

这为我提供了一个高-宽通道排序，它便于数据扩充和预处理功能，如

tf.image.per\u image\u标准化

但是，如果我加载多个图像并希望将它们堆叠在一起（例如，具有多个RGB输入的CNN或多标签语义分割问题），我似乎总是必须复制数据。这里有一种在

tf.stack

中使用副本的方法：

images = []
for image_id in range(0, images):
    file = file_names[image_id]
    images.append(tf.image.decode_png(tf.read_file(file), channels=1)[:, :, 0])
images = tf.stack(images, axis=2)  # Packs as h,w,c

另一种方法是使用

tf.map\fn

，看起来它就是为了这个目的。但是，它“堆叠”在错误的维度，因此我仍然需要一个昂贵的转置：

map = tf.map_fn(lambda f: tf.image.decode_png(tf.read_file(f), channels=1)[:, :, 0],
                file_names, back_prop=False, dtype=tf.uint8)
images = tf.transpose(map, [1, 2, 0])  # from c,h,w to h,w,c

---

是否有可能避免

tf.stack

或

tf.transpose

？

一般来说，删除副本非常困难，因为张量通常是不可变的。每当一个op想要输出一些东西时，它就会分配新的内存并写入

我们可以想象将映射实现更改为沿不同维度堆叠张量。不幸的是，它是使用TensorArray构建的，TensorArray不支持这一点

需要注意的一点是，CHW通常更适合GPU，因为他们更喜欢内部尺寸更大。大多数TF操作都支持这种布局

如果你觉得有冒险精神，你可以试着把这个部分贯穿其中。因为XLA获得了图形的全局视图，所以它可以潜在地优化这些操作中的一些。它正在进行大量开发，可能对您的用例有好处，也可能没有好处

---

您还可以查看使用图像（例如）的官方tensorflow模型以获得最佳实践。

在这种情况下，这是tensorflow数据集API设计中的一个重大缺陷，因为这种简单的转置会给CPU的DRAM带来压力，使我无法最大限度地利用GPU。您所说的NCHW/NHWC在这一点上是不相关的，因为我仍然在预处理所在的Dataset API中，不幸的是，它有某些要求。如果我可以将转置从Dataset API中移出，那么我可以在GPU上运行转置，事情就解决了，但我不能用当前的设计。这不是设计缺陷，而是工程资源分配问题。人们正致力于允许部分数据集API处理在GPU上运行。然而，在这种情况下，如果转置后有更多的CPU预处理要做，那么在GPU上运行转置将没有帮助，因为将数据复制到GPU、执行转置和复制回几乎肯定比在CPU上转置更昂贵。