将预定义函数应用于张量流中的批处理

Question

我有一个用 tensorflow 编写的函数，它可以操作数组。它接收二维数组并对数组的一部分执行一些操作，然后输出一个大小和形状相同但元素经过修改的数组。我的问题是，在我的模型中，我一次不只得到一个数组，而是得到一批形状（batch_size、array_size、channel）。由于通道通常是单声道的，因此输入批次是 4D 张量。 有什么方法可以将我的函数应用于这个 4D 张量中的每一层（矩阵），而无需显着修改函数？

函数类似于：

def myFunction(matrix):
     # function body
     return another_matrix

我如何将它应用于一批 5 个这样的矩阵？

Answer 1

tf.map_fn

它将f 应用于每个元素，在您的情况下应用于每个二维数组。这些二维数组是通过解包第二个参数得到的。