具有多个参数和一个返回值的自定义函数在 map_fn 中用于 Tensorflow 中的张量对象

如何解决具有多个参数和一个返回值的自定义函数在 map_fn 中用于 Tensorflow 中的张量对象

我有两个张量 t1 和 t2（shape=(64,64,3)，dtype=tf.float64）。我想执行一个自定义函数“func”，它以两个张量作为输入并返回一个新的张量。

@tf.function
def func(a):
  t1 = a[0]
  t2 = a[1]
  
  return tf.add(t1,t2)

我正在使用 tensorflow 的 map_fn 为输入的每个元素执行函数。

t = tf.map_fn(fn=func,elems=(t1,t2),dtype=(tf.float64,tf.float64))
tf.print(t)

用于测试目的的样本输入张量是，

t1 = tf.constant([[1.1,2.2,3.3],[4.4,5.5,6.6]])
t2 = tf.constant([[7.7,8.8,9.9],[10.1,11.11,12.12]])

我不能使用带有两个参数的 map_fn。 [尝试使用 tf.stack，也使用 unstack，但这也不起作用。] 知道怎么做吗？

解决方法

“map_fn”的“elems”参数将沿轴 0 传递给它的参数解包。因此，为了在自定义函数中传递多个张量，

1. 我们必须将它们堆叠在一起。
2. 沿轴 0 添加一个额外的维度。

# t1 and t2 has shape [2,3]
val = tf.stack([t1,t2]) # shape is now [2,2,3]
val = tf.expand_dims(val,axis=0) # shape is now [1,3]
t = tf.map_fn(fn=func,elems=val,dtype=tf.float64)

“map_fn”的“dtype”也应该是函数的返回类型。例如，在这种情况下，它应该是 tf.float64。如果函数将返回一个元组，则 dtype 也将是一个元组。

@tf.function
def func(a): # a has shape [2,3]
  t1 = a[0] # shape [2,3]
  t2 = a[1] # shape [2,3]

  return tf.add(t1,t2)