Tensorflow tf.map_fn 参数

Question

我正在尝试构造我的参数，以便它们可以与 tf.map_fn() 一起正常工作，但大多数示例文档仅讨论与函数参数形状相同的数组或张量。

链接包括:

Does tensorflow map_fn support taking more than one tensor?

我的具体例子是这样的:我有一些 tensorflow 函数，需要 [None, 2] 和 [x,y] 作为参数张量形状。

张量 A 的形状为 [batch_size, x*y, 2]

张量 B 的形状为 [batch_size, x, y]

lambdaData = (tensorA, tensorB)
lambdaFunc = lambda x: tensorflowFunc(x[0], x[1])
returnValues = tf.map_fn(lambdaFunc, lambdaData)

来自 tensorflow 文档:

If elems is a (possibly nested) list or tuple of tensors, then each of these 
tensors must have a matching first (unpack) dimension

由于张量A和B只在0维匹配，我不能堆叠或连接它们；我还尝试将 lambdaData 创建为:

1. 两个张量的列表
2. 两个张量的元组
3. 张量对列表

以上所有都会导致不同的维度不匹配错误。我会按照将所有数据放入单个张量的文档中的建议使用，但由于 tensorA 和 tensorB 之间的维度不匹配，我无法做到。有没有人对元素的元组或参数列表有任何好运？

Answer 1

一个可能更漂亮的解决方案是指定 map_fn 的 dtype 参数(参见 documentation )，例如:

tf.map_fn(lambda x: fn(*x), elements, dtype=tf.float32)

如果 fn 只返回一个 float32 值。