把np.einsum翻译成更有表现力的东西

Question

使用python/numpy，我有以下np.einsum

np.einsum('abde,abc->bcde', X, Y)

Y是稀疏的:对于每个[a,b]，只有一个c == 1；所有其他的:= 0。对于轴的相对大小的例子，X.shape是按(1000, 5, 30, 30)的顺序，而Y.shape是等效的(1000, 5, 300)。

这个操作非常昂贵，我想让它更有表现力。首先，einsum不是并行化的。另一方面，由于Y是稀疏的，我有效地计算了我应该执行的乘法运算数的300倍。事实上，当我用n上的循环来写等效的einsum时，我得到了大约3倍的速度。但这显然不是很好。

我该怎么做才能使这个更有表现力呢？我尝试过使用np.tensordot，但是我想不出如何从它得到我想要的东西(我仍然遇到了稀疏/密集的问题)。

Answer 1

如果Y只包含1和0，那么einsum基本上是这样做的：

result = np.zeros(Y.shape[1:] + X.shape[2:], X.dtype)
I, J, K = np.nonzero(Y)
result[J, K] += X[I, J]

但这没有给出正确的结果，因为重复的j，k指数。我无法让numpy.add.at工作，但是只循环这些索引仍然非常快，至少对于给定的形状和稀疏性是这样。

result = np.zeros(Y.shape[1:] + X.shape[2:], X.dtype)
for i, j, k in zip(*np.nonzero(Y)):
    result[j, k] += X[i, j]

这是我使用的测试代码：

a, b, c, d, e = 1000, 5, 300, 30, 30
X = np.random.randint(10, size=(a,b,d,e))
R = np.random.rand(a, b, c)
K = np.argmax(R, axis=2)
I, J = np.indices((a, b), sparse=True)
Y = np.zeros((a, b, c), int)
Y[I, J, K] = 1

Answer 2

用Numba可以很容易地做到这一点。

import numba

@numba.njit('float64[:,:,:,::1](float64[:,:,:,::1], float64[:,:,::1])', fastmath=True, parallel=True)
def compute(x, y):
    na, nb, nd, ne = x.shape
    nc = y.shape[2]
    assert y.shape == (na, nb, nc)

    out = np.zeros((nb, nc, nd, ne))

    for b in numba.prange(nb):
        for a in range(na):
            for c in range(nc):
                yVal = y[a, b, c]
                if np.abs(yVal) != 0:
                    for d in range(nd):
                        for e in range(ne):
                            out[b, c, d, e] += x[a, b, d, e] * yVal

    return out

请注意，在a和b上迭代顺序代码更快。也就是说，为了使代码是并行的，已经交换了循环，并在b上执行并行化(这是一个小轴)。在轴a上进行并行缩减会更有效，但不幸的是，这对于Numba来说并不容易(因为创建线程本地矩阵没有简单的方法，因此需要将矩阵分割成多个块)。

注意，可以用实际值替换nd和ne之类的值(即。30)，因此编译器可以为这个矩阵大小生成更快的代码。

下面是测试代码：

np.random.seed(0)
x = np.random.rand(1000, 5, 30, 30)
y = np.random.rand(1000, 5, 300)
y[np.random.rand(*y.shape) > 0.1] = 0.0 # Make it sparse (90% of 0)

%time res = np.einsum('abde,abc->bcde', x, y)  # 2.350 s
%time res2 = compute(x, y)                     # 0.074 s (0.061 s with hand-written sizes)
print(np.allclose(res, res2))

这是大约32倍的，在一个10核心的英特尔Skylake处理器的。它达到38倍的速度与手写的大小。由于b轴上的并行化，它不能很好地缩放，但是使用其他轴会导致内存访问效率降低。

如果这还不够，那么最好首先转置x和y，这样可以提高数据的局部性(这要归功于a轴上更连续的访问模式)和更好的缩放(通过并行化b和c轴)。尽管如此，转移通常是昂贵的，所以我们当然需要优化它，以获得一个更好的速度。