numpy 点积和矩阵积

Question

我正在使用形状为 (N,)、(N,3) 和 (N,3,3) 的 numpy 数组，它们表示 3D 空间中的标量、向量和矩阵序列。我已经实现了逐点点积、矩阵乘法和矩阵/向量乘法，如下所示：

def dot_product(v, w):
    return np.einsum('ij, ij -> i', v, w)

def matrix_vector_product(M, v):
    return np.einsum('ijk, ik -> ij', M, v)

def matrix_matrix_product(A, B):
    return np.einsum('ijk, ikl -> ijl', A, B)

如您所见，我使用 einsum 是因为没有更好的解决方案。令我惊讶的是，我无法使用 np.dot... 这似乎不适合这种需要。有没有更 numpythonic 的方式来实现这些功能？

特别是如果函数可以通过广播第一个缺失的轴也可以在形状 (3,) 和 (3,3) 上工作，那就太好了。我想我需要省略号，但我不太明白如何实现结果。

Answer 1

这些操作不能重构成一般的 BLAS 调用，对于这种大小的数组，循环 BLAS 调用会非常慢。因此，einsum 可能是此类操作的最佳选择。

您的函数可以用省略号概括如下：

def dot_product(v, w):
    return np.einsum('...j,...j->...', v, w)

def matrix_vector_product(M, v):
    return np.einsum('...jk,...k->...j', M, v)

def matrix_matrix_product(A, B):
    return np.einsum('...jk,...kl->...jl', A, B)

Answer 2

就像工作笔记一样，这 3 个计算也可以写成：

np.einsum(A,[0,1,2],B,[0,2,3],[0,1,3])
np.einsum(M,[0,1,2],v,[0,2],[0,1]) 
np.einsum(w,[0,1],v,[0,1],[0])

或者用 Ophion 的概括

np.einsum(A,[Ellipsis,1,2], B, ...)

根据输入数组的维度生成[0,1,..] 列表应该不难。

---

通过专注于泛化 einsum 表达式，我错过了您试图重现的是 N 小点积的事实。

np.array([np.dot(i,j) for i,j in zip(a,b)])

值得记住的是np.dot 使用快速编译的代码，并且专注于数组很大的计算。您的问题是计算许多小点积之一。

在没有定义轴的额外参数的情况下，np.dot 仅执行两种可能的组合，可以表示为：

np.einsum('i,i', v1, v2)
np.einsum('...ij,...jk->...ik', m1, m2)

dot 的运算符版本将面临同样的限制 - 没有额外的参数来指定如何组合轴。

注意tensordot 对概括dot 所做的工作也可能具有指导意义：

def tensordot(a, b, axes=2):
    ....
    newshape_a = (-1, N2)
    ...
    newshape_b = (N2, -1)
    ....
    at = a.transpose(newaxes_a).reshape(newshape_a)
    bt = b.transpose(newaxes_b).reshape(newshape_b)
    res = dot(at, bt)
    return res.reshape(olda + oldb)

它可以在多个轴上执行dot 求和。但是转置和整形完成后，计算就变成了标准的dot 2d数组。

---

这可能已被标记为重复问题。一段时间以来，人们一直在询问是否要做多个点积。

Matrix vector multiplication along array axes 建议使用numpy.core.umath_tests.matrix_multiply

https://stackoverflow.com/a/24174347/901925 等于：

matrix_multiply(matrices, vectors[..., None])
np.einsum('ijk,ik->ij', matrices, vectors)

matrix_multiply 的 C 文档说明：

* This implements the function
* out[k, m, p] = sum_n { in1[k, m, n] * in2[k, n, p] }.

来自同一目录的inner1d 对(N,n) 向量执行相同操作

inner1d(vector, vector)  
np.einsum('ij,ij->i', vector, vector)
# out[n] = sum_i { in1[n, i] * in2[n, i] }

两者都是UFunc，可以处理最右侧维度的广播。在numpy/core/test/test_ufunc.py 中，这些函数用于行使UFunc 机制。

matrix_multiply(np.ones((4,5,6,2,3)),np.ones((3,2)))

https://stackoverflow.com/a/16704079/901925补充说这种计算可以用*和sum来完成，例如

(w*v).sum(-1)
(M*v[...,None]).sum(-1)
(A*B.swapaxes(...)).sum(-1)

---

在进一步的测试中，我认为inner1d 和matrix_multiply 匹配您的dot 和matrix-matrix 产品案例，如果您添加[...,None]，则matrix-vector 案例。看起来它们比 einsum 版本快 2 倍（在我的机器和测试阵列上）。

https://github.com/numpy/numpy/blob/master/doc/neps/return-of-revenge-of-matmul-pep.rst 是对numpy 上的@ 中缀运算符的讨论。我认为numpy 开发人员对这个 PEP 的热情不如 Python 开发人员。