Question

我目前正在编写代码，我需要尽可能快地计算三种二维数组之间的内在产品。

我们称他们为a，b，c。它们都具有相同的尺寸（N×M）。

我想计算以下3-d数组op，大小（N x N x N），这样op [i，j，k]是a [i，m] b的m之和[j，m] c [k，m]

这基本上是np.inner到3个输入的扩展版本，而不是2。

在实践中，我将遇到的维度类似于N = 100和M = 300 000.矩阵根本不会稀疏，因此op包含大约100万个非零值。

到目前为止，我尝试了两种方法。第一个使用广播：

import numpy as np
N = 100
M = 300000
a = np.random.randn(N, M)
b = np.random.randn(N, M)
c = np.random.randn(N, M)

def method1(a, b, c):
    a_i = a[:, None, None, :]
    b_j = b[None, :, None, :]
    c_k = c[None, None, :, :]
    return np.sum(a_i * b_j * c_k, axis=3)

这个问题在于它首先计算a_i * b_j * c_k这是一个N x N x N x M数组，所以在我的情况下它只是处理得太多了。

我尝试过使用np.einsum的另一种方法，它比上一种方法快得多：

def method2(a, b, c):
    return np.einsum('im,jm,km', a, b, c)

我的问题是它仍然太慢。对于N = 100和M = 30 000，在我的计算机上运行已经需要95秒，因此将M设为其实际值300 000是不可能的。

我的问题是：你知道解决我的问题的任何pythonic方式（也许是一个神奇的numpy函数吗？），或者我是否必须求助于cython或numba这样的事实才能使这个计算变得可行？

提前感谢您的帮助！

Answer 1

非常有趣，与this other problem有关。

方法＃1：对于体面的大小数组

基于上述Q&A的获胜方法，这里有一个解决方案 -

np.tensordot(a[:,None]*b,c,axes=(2,1))

说明：

1）a[:,None]*b：获取形状(N, N, M)的3D数组。因此，对于用例，它将是(100, 100, 30000)，对于常规系统来说可能有点太多了，但是可能只是给出一些额外的系统内存 juice 。

2）np.tensordot(..)：接下来，我们将使用张量点对第三个数组c求和 - 减少前一步的最后一个轴，使其具有(100, 100, 100)形状的输出数组。

方法＃2：对于非常大的数组，b与c相同

out = np.zeros((N, N, N))
for i in range(N):
    for j in range(N):
        for k in range(j+1):
            out[i,j,k] = np.einsum('i,i,i->',a[i],b[j],b[k])

r,c = np.triu_indices(N,1)
out[np.arange(N)[:,None], r,c] = out[np.arange(N)[:,None], c,r]

python中两个以上矩阵的快速内积

1 个答案: