tensorflow - 我可以将 tf.map_fn(...) 应用于多个输入/输出吗？

Question

a = tf.constant([[1,2,3],[4,5,6]])
b = tf.constant([True, False], dtype=tf.bool)

a.eval()
array([[1, 2, 3],
       [4, 5, 6]], dtype=int32)
b.eval()
array([ True, False], dtype=bool)

我想对上面的输入应用一个函数， a , 和 b使用 tf.map_fn .它将同时输入 [1,2,3] , 和 True并输出相似的值。

假设 out 函数只是身份: lambda(x,y): x,y所以，给定输入 [1,2,3], True ，它将输出那些相同的张量。

我知道如何使用 tf.map_fn(...)一个变量，但不是两个。在这种情况下，我有混合数据类型(int32 和 bool)，所以我不能简单地连接张量并在调用后拆分它们。

我可以用 tf.map_fn(...)具有不同数据类型的多个输入/输出？

Answer 1

弄清楚了。您必须在 dtype 中为每个张量定义数据类型对于每个不同的张量，您可以将张量作为元组传递，您的 map 函数接收一个输入元组，然后 map_fn返回一个元组。

有效的例子:

a = tf.constant([[1,2,3],[4,5,6]])
b = tf.constant([True, False], dtype=tf.bool)

c = tf.map_fn(lambda x: (x[0], x[1]), (a,b), dtype=(tf.int32, tf.bool))

c[0].eval()
array([[1, 2, 3],
       [4, 5, 6]], dtype=int32)
c[1].eval()
array([ True, False], dtype=bool)