Question

我坚持将两个嵌套的 for 循环向量化以计算两个矩阵的内积。这是一个代码示例：

a = tf.random.normal((16, 128, 128, 128))
b = tf.random.normal((16, 128, 128, 128))

out = []
for n in range(16):
    for l in range(n + 1):
        out.append(tf.einsum("bjk, bj -> b", a[l], b[n-l, l]))
tf.stack(out)

到目前为止我尝试过的：

a_list = []
b_list = []
for n in range(16):
    for l in range(n+1):
        a_list.append([l])
        b_list.append([n-l, l])   
tf.einsum("xbjk, xbj -> xb", tf.gather_nd(a, a_list), tf.gather_nd(b, b_list))

这有效，但不会导致计算速度加快。当我将第一个作为 @tf.function 运行时，我得到 29.9 毫秒（没有 76 毫秒），而使用 tf.gather_nd 的一个需要 235 毫秒。

是否有更好的方法来向量化计算，或者 @tf.function 可能使用 tf.while_loop 进行此计算的最佳方法？

使用 Tensorflow 向量化嵌套的 for 循环

0 个答案: