Question

在TensorFlow中，我有一个形状为M的2级张量[D, D]（矩阵）和一个形状为T的3级张量[D, D, D]。

我需要将它们组合成一个新的矩阵R，如下所示：元素R[a, b+c-a]由所有元素T[a, b, c]*M[b, c]的总和给出，其中b+c-a是常数（其中b+c-a必须在0到D-1之间）。

创建R的一种无效方法是使用嵌套的for遍历索引，并检查b+c-a是否超过D-1（例如numpy）：

R = np.zeros([D,D])

for a in range(D):
  for b in range(D):
    for c in range(D):
      if 0 <= b+c-a < D:
        R[a, b+c-a] += T[a, b, c]*M[b, c]

但是我想使用广播和/或其他更有效的方法。

我该如何实现？

Answer 1

您可以按如下方式向量化该计算：

import numpy as np

np.random.seed(0)
D = 10
M = np.random.rand(D, D)
T = np.random.rand(D, D, D)
# Original calculation
R = np.zeros([D, D])
for a in range(D):
    for b in range(D):
        for c in range(D):
            if 0 <= b + c - a < D:
                R[a, b + c - a] += T[a, b, c] * M[b, c]
# Vectorized calculation
tm = T * M
a = np.arange(D)[:, np.newaxis, np.newaxis]
b, c = np.ogrid[:D, :D]
col_idx = b + c - a
m = (col_idx >= 0) & (col_idx < D)
row_idx = np.tile(a, [1, D, D])
R2 = np.zeros([D, D])
np.add.at(R2, (row_idx[m], col_idx[m]), tm[m])
# Check result
print(np.allclose(R, R2))
# True

或者，您可以考虑使用Numba来加速循环：

import numpy as np
import numba as nb

@nb.njit
def calculation_nb(T, M, D):
    tm = T * M
    R = np.zeros((D, D), dtype=tm.dtype)
    for a in nb.prange(D):
      for b in range(D):
        for c in range(max(a - b, 0), min(D + a - b, D)):
          R[a, b + c - a] += tm[a, b, c]
    return R

print(np.allclose(R, calculation_nb(T, M, D)))
# True

在一些快速测试中，即使没有并行化，这也比NumPy快得多。

TensorFlow（或Numpy）中的高级广播

1 个答案: