python - 使用 tf.map_fn 将多个图像读取为张量

Question

我正在使用 Tensorflow 的数据集 API 读取各种图像(数据和标签)。由于数据集队列在 CPU 上，因此复制数据的成本很高。但是，我似乎找不到避免这种情况的方法。

问题:我能否以统一的顺序(例如 h、w、c)高效地加载各种图像？

假设我想读取单个单 channel 图像，我可以按如下方式进行:

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

或者对于多 channel RGB:

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

这给了我一个高度-宽度- channel 顺序，这对于数据增强和预处理功能很方便，例如 tf.image.per_image_standardization。

但是，如果我加载多张图像并想将它们堆叠在一起(例如，具有多个 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？

Answer 1

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

可以想象更改 map_fn 实现以沿不同维度堆叠张量。不幸的是，它是使用不支持此功能的 TensorArray 构建的。

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

如果你喜欢冒险，你可以尝试通过 XLA 运行这部分.因为 XLA 获得了图形的全局 View ，所以它可以潜在地优化其中的一些操作。它正在大力开发中，可能会也可能不会有益于您的用例。

您还可以查看处理图像的官方 tensorflow 模型(例如 https://github.com/tensorflow/models/tree/master/official/resnet)以获得最佳实践。