2向量矩阵点积的性能

Performance of 2 vector-matrix dot product

问题更侧重于计算性能。

我有 2 个矢量矩阵。这意味着它们对于 X、Y、Z 具有 3 个深度维度。矩阵的每个元素都必须与另一个矩阵相同位置的元素进行点积。

一个简单而低效的代码将是这个:

import numpy as np
a = np.random.uniform(low=-1.0, high=1.0, size=(1000,1000,3))
b = np.random.uniform(low=-1.0, high=1.0, size=(1000,1000,3))
c = np.zeros((1000,1000))
numRow,numCol,numDepth = np.shape(a)

for idRow in range(numRow):
    for idCol in range(numCol):
        # Angle in radians
        c[idRow,idCol] = math.acos(a[idRow,idCol,0]*b[idRow,idCol,0] + a[idRow,idCol,1]*b[idRow,idCol,1] + a[idRow,idCol,2]*b[idRow,idCol,2])

然而,numpy 函数可以加快计算速度,如下所示,使代码更快:

# Angle in radians
d = np.arccos(np.multiply(a[:,:,0],b[:,:,0]) + np.multiply(a[:,:,1],b[:,:,1]) + np.multiply(a[:,:,2],b[:,:,2]))

但是,我想知道是否有其他的 sintaxis 可以改进上面的这个,也许还有其他的功能,索引,...

第一个代码需要 4.658s 而第二个需要 0.354s

您可以使用 np.einsum 执行此操作,它在任何轴上相乘然后求和:

np.arccos(np.einsum('ijk,ijk->ij', a, b))

完成您在问题中发布的内容的更直接的方法是使用 np.sum,其中您沿最后一个轴求和 (-1):

np.arccos(np.sum(a*b, -1))

他们都给出相同的答案,但 einsum 最快,sum 其次:

In [36]: timeit np.arccos(np.einsum('ijk,ijk->ij', a, b))
10000 loops, best of 3: 20.4 µs per loop

In [37]: timeit e = np.arccos(np.sum(a*b, -1))
10000 loops, best of 3: 29.8 µs per loop

In [38]: %%timeit
   ....: d = np.arccos(np.multiply(a[:,:,0],b[:,:,0]) +
   ....:               np.multiply(a[:,:,1],b[:,:,1]) +
   ....:               np.multiply(a[:,:,2],b[:,:,2]))
   ....:
10000 loops, best of 3: 34.6 µs per loop

Pythran 编译器可以通过以下方式进一步优化您的原始表达式:

  • 删除临时数组
  • 使用 SIMD 指令
  • 使用多线程

如本例所示:

$ cat cross.py
#pythran export cross(float[][][], float[][][])
import numpy as np
def cross(a,b):
    return np.arccos(np.multiply(a[:, :, 0], b[:, :, 0]) + np.multiply(a[:, :, 1],b[:, :, 1]) + np.multiply(a[:, :, 2], b[:, :, 2]))
$ python -m timeit -s 'import numpy as np; a = np.random.uniform(low=-1.0, high=1.0, size=(1000, 1000, 3)); b = np.random.uniform(low=-1.0, high=1.0, size=(1000, 1000, 3)); c = np.zeros((1000, 1000)); from cross import cross' 'cross(a,b)'
10 loops, best of 3: 35.4 msec per loop
$ pythran cross.py -DUSE_BOOST_SIMD -fopenmp -march=native
$ python -m timeit -s 'import numpy as np; a = np.random.uniform(low=-1.0, high=1.0, size=(1000, 1000, 3)); b = np.random.uniform(low=-1.0, high=1.0, size=(1000, 1000, 3)); c = np.zeros((1000, 1000)); from cross import cross' 'cross(a,b)'
100 loops, best of 3: 11.8 msec per loop