python - python中巨大矩阵的点积的行总和

Question

我有 2 个矩阵 100kx200 和 200x100k如果它们是小矩阵，我只会使用 numpy 点积

sum(a.dot(b), axis = 0)

但是矩阵太大了，而且我不能使用循环，有没有一种聪明的方法可以做到这一点？

Answer 1

一个可能的优化是

>>> numpy.sum(a @ b, axis=0)
array([  1.83633615,  18.71643672,  15.26981078, -46.33670382,  13.30276476])
>>> numpy.sum(a, axis=0) @ b
array([  1.83633615,  18.71643672,  15.26981078, -46.33670382,  13.30276476])

计算 a @ b 需要 10k×200×10k 次操作，而首先对行求和会将乘法减少到 1×200×10k 次操作，提高 10k× 。

这主要是因为认识

   numpy.sum(x, axis=0) == [1, 1, ..., 1] @ x
=> numpy.sum(a @ b, axis=0) == [1, 1, ..., 1] @ (a @ b)
                            == ([1, 1, ..., 1] @ a) @ b
                            == numpy.sum(a, axis=0) @ b

其他轴类似。

>>> numpy.sum(a @ b, axis=1)
array([  2.8794171 ,   9.12128399,  14.52009991,  -8.70177811, -15.0303783 ])
>>> a @ numpy.sum(b, axis=1)
array([  2.8794171 ,   9.12128399,  14.52009991,  -8.70177811, -15.0303783 ])

(注意:x @ y 相当于 x.dot(y) for 2D matrixes and 1D vectors on Python 3.5+ with numpy 1.10.0+ )

---

$ INITIALIZATION='import numpy;numpy.random.seed(0);a=numpy.random.randn(1000,200);b=numpy.random.rand(200,1000)'

$ python3 -m timeit -s "$INITIALIZATION" 'numpy.einsum("ij,jk->k", a, b)'
10 loops, best of 3: 87.2 msec per loop

$ python3 -m timeit -s "$INITIALIZATION" 'numpy.sum(a@b, axis=0)'
100 loops, best of 3: 12.8 msec per loop

$ python3 -m timeit -s "$INITIALIZATION" 'numpy.sum(a, axis=0)@b'
1000 loops, best of 3: 300 usec per loop

---

插图:

In [235]: a = np.random.rand(3,3)
array([[ 0.465,  0.758,  0.641],
       [ 0.897,  0.673,  0.742],
       [ 0.763,  0.274,  0.485]])

In [237]: b = np.random.rand(3,2)
array([[ 0.303,  0.378],
       [ 0.039,  0.095],
       [ 0.192,  0.668]])

现在，如果我们简单地执行 a @ b，我们将需要 18 个乘法运算和 6 个加法运算。另一方面，如果我们执行 np.sum(a, axis=0) @b，我们只需要 6 个乘法运算和 2 个加法运算。改进了 3 倍，因为我们在 a 中有 3 行。至于 OP 的情况，这应该比简单的 a @b 计算提高 10k 倍，因为他在 a 中有 10k 行。