制作自定义（tf.function）函数二维数组

Question

我正在尝试将变量 x、y 的函数的二维数组计算为 tf.function。该函数相当复杂，我想制作该函数的二维数组，其中 x 和 y 采用值列表（tf.linspace）。我试过输入这样一个函数的相关参数，这是它的样子

@tf.function
def function_matrix(xi, xf, yi, yf, num , some_other_args):
    
    #part1
    M=np.zeros((num, num))
    xlist=tf.linspace(xi, xf, num)
    ylist=tf.linspace(yi, yf, num)
    
    #part2
    for x in range(num):
         for y in range(num):
             M[x,y]=some_complicated_function(xlist[x], ylist[y], some_other_args)     #this is also a @tf.function
    
    return (M)

我遇到的问题是，在 tf.function 中，如果我尝试访问像 xlist[x] 这样的数组元素，结果是 Tensor("strided_slice:0", shape=(), dtype=float64)。所以当在 some_complicated_function 中传递这个值时，我得到一个错误“设置一个带有序列的数组元素”。如果 function_matrix 不是 tf.function，则不会发生此类错误。有人可以帮忙吗？至于我哪里可能出错？或者我可以计算相当复杂函数的二维矩阵的任何替代方法？ 任何帮助将不胜感激，谢谢！

我的尝试： 第 1 部分运行良好，如果我将 xlist 作为函数的输出返回，我会得到一个普通数组 tf.Tensor( [the_array_here], shape=(num,), dtype=float64)。同样，如果输出是 xlist[index]，我得到tf.Tensor( [the_element_here], shape=(), dtype=float64)。但是我是否尝试从函数中打印 xlist[index]，我得到Tensor("strided_slice:0", shape=(), dtype=float64)。所以我得出的结论是，tf 以某种方式将 xlist[index] 视为某种占位符。但我不知道为什么...

Answer 1

哦，好问题！ tensorflow 真的不喜欢 for 循环，它是 python 无法自动转换为 tensorflow graph representation 的代码。实现这一点的方法是生成要在张量中操作的网格。比方说：

xlist=[1,2] # this is a tf.Tensor
ylist=[1,2] # this is a tf.Tensor

那么，使用tf.meshgrid，你应该构造xylist：

xylist=[[1,1], [1,2], [2,1], [2,2]] # this is a tf.Tensor

然后使用tf.map_fn 将您的函数应用于每一对。

M = tf.map_fn(xylist, some_complicated_function)
M = tf.reshape(M, (...))

注意，如果some_complicated_function 包含任何非tensorflow 代码（或无法自动转换的代码），例如使用numpy、pandas、pillow...，您可以将其包裹在@ 987654337@ - 但现在这种方式违背了将函数转换为tf.function 的目的。 （编辑：我现在看到您说：# this is also a tf.function，这意味着您不必将其包装在 tf.py_function 中）

您还可以通过在xylist 中将extra_args 附加到每一对 来包含extra_args（是的，每一对，即使它们是不变的）。

TL;DR：使用 tf.map_fn 而不是嵌套的 for 循环。