假设我有两个向量并希望计算它们的点积;这很简单,
import numpy as np
a = np.random.rand(3)
b = np.random.rand(3)
result = np.dot(a,b)
如果我有一堆向量,想要对每一个进行点乘,最简单的代码是:
# 5 = number of vectors
a = np.random.rand(5,3)
b = np.random.rand(5,3)
result = [np.dot(aa,bb) for aa, bb in zip(a,b)]
有两种批处理计算的方式,一种是使用乘法和求和,另一种是使用einsum。
result = np.sum(a*b, axis=1)
# or
result = np.einsum('ij,ij->i', a, b)
然而,这两种方法都没有使用到BLAS后端,只能使用单个核心。当N
很大时,比如100万时,效果不是很好。
tensordot确实使用了BLAS后端。使用tensordot进行此计算的一种糟糕方式是:
np.diag(np.tensordot(a,b, axes=[1,1])
这很糟糕,因为它分配了一个 N*N
的矩阵,而大部分元素都是无用的。
另一种(极快速的)方法是使用隐藏的 inner1d 函数。
from numpy.core.umath_tests import inner1d
result = inner1d(a,b)
但是看起来 这种方式似乎不可行,因为可能会公开出现的问题已经得不到解决了。而这仍然需要用C语言编写循环,而不是使用多个核心。
有没有一种方法可以让
dot
、matmul
或tensordot
一次性在多个核心上执行所有这些点积运算呢?
a[:,None,:]@b[:,:,None]
. Also tryeinsum
withoptimize=True
- hpauljN
并更改标志的解决方案。 - Brandon Dubematmul
比einsum
更快。但是在这种情况下,可能是由于批处理大小与“inner”维度之间的比较,使得einsum
具有优势。 - hpaulj