Numpy批量点积

Question

Numpy批量点积

8

假设我有两个向量并希望计算它们的点积；这很简单，

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一次性在多个核心上执行所有这些点积运算呢？

- Brandon Dube

a[:,None,:]@b[:,:,None]. Also try einsum with optimize=True - hpaulj

matmul调用不使用多个核心，比inner1d慢得多（N = 1024，inner1d = 4usec，matmul = 19usec）。优化的einsum有一个巨大的常数时间开销（~25usec），但对于适度大的N（100万）而言，它大约快10％。我更喜欢一种“总是良好”的解决方案，而不是需要嗅探N并更改标志的解决方案。 - Brandon Dube

通常情况下，matmul 比 einsum 更快。但是在这种情况下，可能是由于批处理大小与“inner”维度之间的比较，使得 einsum 具有优势。 - hpaulj

这取决于矩阵或向量的大小。另外，使用一些已知信息（例如向量或矩阵始终具有相同的大小）可以进一步优化。在更大的问题上，并行化也是有益的。小矩阵的例子：https://stackoverflow.com/a/59356461/4045774 - max9111

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Jérôme Richard · Accepted Answer

首先，没有直接的BLAS函数可以完成此操作。使用多个级别1的BLAS函数调用效率不高，因为为非常短暂的计算使用多个线程会引入相当大的开销，而不使用多个线程可能是次优的。然而，这种计算主要是内存绑定的，因此在拥有许多核心的平台上扩展性较差（几个核心通常足以饱和内存带宽）。

一个简单的解决方案是使用Numexpr包，这应该能够有效地实现（它应该避免创建临时数组，并且应该也使用了多个线程）。然而，在这种情况下，对于大型数组，性能有些令人失望。

最好的解决方案似乎是使用Numba（或Cython）。Numba可以为小型和大型输入数组生成快速代码，并且易于并行化代码。但是，请注意，管理线程会引入一些开销，对于小型数组可能相当大(在某些许多核心的平台上高达几毫秒)。

以下是Numexpr的实现：

import numexpr as ne
expr = ne.NumExpr('sum(a * b, axis=1)')
result = expr.run(a, b)

这里是一个（顺序的）Numba实现:

import numba as nb

# Use `parallel=True` for a parallel implementation
@nb.njit('float64[:](float64[:,::1], float64[:,::1])')
def multiDots(a, b):
    assert a.shape == b.shape
    n, m = a.shape
    res = np.empty(n, dtype=np.float64)

    # Use `nb.prange` instead of `range` to run the loop in parallel
    for i in range(n):
        s = 0.0
        for j in range(m):
            s += a[i,j] * b[i,j]
        res[i] = s

    return res

result = multiDots(a, b)

以下是一台（旧的）双核机器上的一些基准测试：

On small 5x3 arrays:
    np.einsum('ij,ij->i', a, b, optimize=True):  45.2 us
    Numba (parallel):                            12.1 us
    np.sum(a*b, axis=1):                          9.5 us
    np.einsum('ij,ij->i', a, b):                  6.5 us
    Numexpr:                                      3.2 us
    inner1d(a, b):                                1.8 us
    Numba (sequential):                           1.3 us

On small 1000000x3 arrays:
    np.sum(a*b, axis=1):                         27.8 ms
    Numexpr:                                     15.3 ms
    np.einsum('ij,ij->i', a, b, optimize=True):   9.0 ms
    np.einsum('ij,ij->i', a, b):                  8.8 ms
    Numba (sequential):                           6.8 ms
    inner1d(a, b):                                6.5 ms
    Numba (parallel):                             5.3 ms

顺序Numba实现提供了良好的权衡。如果您真的想获得最佳性能，可以使用开关。然而，在独立于平台的情况下选择最佳的n阈值并不那么容易。