在 Numpy 中并行化类似 matmul 的矩阵计算答案

【问题标题】：Parallelize a matmul-like matrix computation in Numpy在 Numpy 中并行化类似 matmul 的矩阵计算
【发布时间】：2020-09-04 23:35:52
【问题描述】：

在形状(n,n,m,m)中输入X，

以(n,n) 的形式输出Y，其中Y[i,j]=∑_{k=1}^{n}{||X[i,j]-X[i,k]*X[k,j]||}，* 表示逐点乘法。

愚蠢的for循环版本是这样的：

X = np.random.randint(1,10,size=(5,5,3,3))
n, _, m, _ = X.shape
Y = np.zeros((n, n))
for i in range(n):
    for j in range(n):
        cnt = 0.0
        X_ij = X[i, j] # in shape m x m
        for k in range(n):
            X_ikj = X[i, k] * X[k, j] # point-wise, in shape m x m
            cnt += np.sum(np.abs(X_ij - X_ikj))
        Y[i, j] = cnt

但是我想使用 numpy 并行矩阵计算。恰好Y[i,j]=∑_{k=1}^{n}{||X[i,j]-X[i,k]*X[k,j]||} 与matmul 具有相似的形式。所以在我看来基本上有两点：

如何matmul 仅沿前两个维度，同时保持后两个维度的逐点乘法？
matmul 已经总结了n-dim 向量{X[i,k]*X[k,j]}_{k in [1,n]}。但是，在汇总之前，每个 X[i,k]*X[k,j] 都有一个函数。

感谢任何可能的想法！谢谢。

【问题讨论】：

您能否说明“在汇总之前将函数应用到每个X[i,k]*X[k,j]”的意思以及这一点在问题中的相关性是什么？
@amzon-ex 说f(X,i,k,j)=||X[i,j]-X[i,k]*X[k,j]||。所以不可能对每个X[i,j]直接应用f，然后再应用matmul。
如果 X_ikj = np.matmul(X[i, k],X[k, j]) 而不是逐点，请在 Quang 的迷人代码中使用 pairwise_consistency = 1 - np.abs(X[:, :, None] - np.matmul(X[:, None],X_t[None, ...])).sum((2, 3, 4)) / (2 * n * m)。

标签： arrays numpy matrix matrix-multiplication tensor

【解决方案1】：

你可以使用broadcasting，但你需要用转置交换两个轴：

np.random.seed(1)
X = np.random.randint(1,10,size=(5,5,3,3))

# transpose
# so X_t[j,k] == X[k,j]
X_t = X.transpose(1,0,2,3)

# output
# X_t[None,...]*X[:,None] is X[k,j] * X[i,k]
ret = np.abs(X[:,:,None] - X_t[None,...]*X[:,None]).sum((2,3,4))

# check
(ret==Y).all()
# True

输出 (ret)

array([[1108, 1078,  709,  825,  752],
       [1163, 1185,  988, 1034,  910],
       [1043,  973,  828,  926,  706],
       [ 908,  927,  800, 1078,  765],
       [ 990,  905,  662,  864,  865]])

【讨论】：