计算大量3x3点积的最快方法

Question

我需要计算大量的3x3线性变换（例如旋转）。这是我目前所拥有的：

import numpy as np
from scipy import sparse
from numba import jit

n = 100000 # number of transformations
k = 100 # number of vectors for each transformation

A = np.random.rand(n, 3, k) # vectors
Op = np.random.rand(n, 3, 3) # operators
sOp = sparse.bsr_matrix((Op, np.arange(n), np.arange(n+1))) # same as Op but as block-diag

def dot1():
    """ naive approach: many times np.dot """
    return np.stack([np.dot(o, a) for o, a in zip(Op, A)])

@jit(nopython=True)
def dot2():
    """ same as above, but jitted """
    new = np.empty_like(A)
    for i in range(Op.shape[0]):
        new[i] = np.dot(Op[i], A[i])
    return new

def dot3():
    """ using einsum """
    return np.einsum("ijk,ikl->ijl", Op, A)

def dot4():
    """ using sparse block diag matrix """
    return sOp.dot(A.reshape(3 * n, -1)).reshape(n, 3, -1)

在 MacBook Pro 2012 上，这给了我：

In [62]: %timeit dot1()
783 ms ± 20.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [63]: %timeit dot2()
261 ms ± 1.93 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [64]: %timeit dot3()
293 ms ± 2.89 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [65]: %timeit dot4()
281 ms ± 6.15 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

除了朴素方法外，其他方法都很相似。有没有一种方法可以显著加速这个过程？

编辑

(如果可用，cuda方法是最佳的。以下是非cuda版本的比较)

根据各种建议，我修改了dot2，添加了Op@A方法以及基于#59356461的一个版本。

@njit(fastmath=True, parallel=True)
def dot2(Op, A):
    """ same as above, but jitted """
    new = np.empty_like(A)
    for i in prange(Op.shape[0]):
        new[i] = np.dot(Op[i], A[i])
    return new

def dot5(Op, A):
    """ using matmul """
    return Op@A

@njit(fastmath=True, parallel=True)
def dot6(Op, A):
    """ another numba.jit with parallel (based on #59356461) """
    new = np.empty_like(A)
    for i_n in prange(A.shape[0]):
        for i_k in range(A.shape[2]):
            for i_x in range(3):
                acc = 0.0j
                for i_y in range(3):
                    acc += Op[i_n, i_x, i_y] * A[i_n, i_y, i_k]
                new[i_n, i_x, i_k] = acc
    return new

这是在另一台机器上使用benchit得到的结果：

def gen(n, k):
    Op = np.random.rand(n, 3, 3) + 1j * np.random.rand(n, 3, 3)
    A = np.random.rand(n, 3, k) + 1j * np.random.rand(n, 3, k)
    return Op, A

# benchit
import benchit
funcs = [dot1, dot2, dot3, dot4, dot5, dot6]
inputs = {n: gen(n, 100) for n in [100,1000,10000,100000,1000000]}

t = benchit.timings(funcs, inputs, multivar=True, input_name='Number of operators')
t.plot(logy=True, logx=True)

- piliv

---

你试过 Op@A 吗？ - hpaulj

你可以使用不同的数据类型吗？需要什么样的精度？ - Ehsan

@hpaulj 没有，谢谢！ @Ehsan 实际上，我需要复杂的。我没有想到这会有太大的区别，所以在我的问题中使用了浮点数。结果发现 einsum 在处理复杂数据时表现得更差。 - piliv

这个 https://stackoverflow.com/a/59356461/4045774 在(3x3)x(3x100)的情况下运行速度大约是op@A的两倍。同时请记住，这个问题在很大程度上受到内存带宽的限制。如果你在GPU上计算它，还要考虑从GPU复制数据的时间。 - max9111

Answer 1

你已经得到了一些很好的建议，但是由于这个具体的目标，我想再添加一个：

“是否有一种方法可以显著加速这个过程？”

实际上，如果你需要这些操作变得“显著”快速（通常意味着> 10倍），你可能会希望使用GPU来进行矩阵乘法运算。以下是一个快速的示例：

import numpy as np
import cupy as cp

n = 100000 # number of transformations
k = 100 # number of vectors for each transformation

# CPU version
A = np.random.rand(n, 3, k) # vectors
Op = np.random.rand(n, 3, 3) # operators

def dot5(): # the suggested, best CPU approach
    return Op@A


# GPU version using a V100
gA = cp.asarray(A)
gOp = cp.asarray(Op)

# run once to ignore JIT overhead before benchmarking
gOp@gA;

%timeit dot5()
%timeit gOp@gA; cp.cuda.Device().synchronize() # need to sync for a fair benchmark
112 ms ± 546 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
1.19 ms ± 1.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

- Nick Becker

Answer 2

按照评论中@hpaulj的建议，使用Op@A。

下面是使用benchit进行比较：

def dot1(A,Op):
    """ naive approach: many times np.dot """
    return np.stack([np.dot(o, a) for o, a in zip(Op, A)])

@jit(nopython=True)
def dot2(A,Op):
    """ same as above, but jitted """
    new = np.empty_like(A)
    for i in range(Op.shape[0]):
        new[i] = np.dot(Op[i], A[i])
    return new

def dot3(A,Op):
    """ using einsum """
    return np.einsum("ijk,ikl->ijl", Op, A)

def dot4(A,Op):
    n = A.shape[0]
    sOp = sparse.bsr_matrix((Op, np.arange(n), np.arange(n+1))) # same as Op but as block-diag
    """ using sparse block diag matrix """
    return sOp.dot(A.reshape(3 * n, -1)).reshape(n, 3, -1)

def dot5(A,Op):
  return Op@A

in_ = {n:[np.random.rand(n, 3, k), np.random.rand(n, 3, 3)] for n in [100,1000,10000,100000,1000000]}

对于更大规模的应用，它们的性能看起来非常接近，但是dot5略微更快。

- Ehsan

---

谢谢！这是否意味着dot2...dot5之间的区别主要在于开销？（数组初始化，复制，JIT等） - piliv

@piliv 我不是完全确定实现的方法，但是如果要猜测，我会说是的。 - Ehsan

Answer 3

在一个回答中，尼克提到使用GPU - 当然这是最好的解决方案。但通常情况下，你所做的可能是CPU受限的。因此（除了GPU方法），你可以获得的最佳性能是利用机器上的所有核心并行工作。为此，你需要使用multiprocessing（而不是Python的多线程！）将任务分割成一块一块的，在每个核心上并行运行。虽然这不是微不足道的事情，但其实也不太难，网上有许多好的例子和指南。但如果你有一台8核机器，它可能会给你近乎8倍的速度增加，只要你小心，避免通过尝试在进程之间传递许多小对象来引起内存瓶颈，但可以在开始时将它们都传递给一组。

- Richard

---

这个例子实际上受到了内存带宽的限制，因为计算工作太简单了，无法让所有计算单元保持繁忙状态。即使是分配输出数组，也需要占用一半的运行时间。 - max9111

...这就是为什么我在我写的内容中放上了斜体字 :D 实际上，只要在执行之前将数组分组并传递给每个子进程，您可能仍然可以使用MP获得良好的效果。 ...你说得对，对每个单独的数组进行简单的MP类型map或等效操作，可能会遭受过多的内存洗牌。 - Richard

谢谢，但是numpy不是已经在使用多进程进行点积/矩阵乘法了吗？（我正在使用conda的numpy） - piliv

我不确定 - 可能是。对于你的情况，它可能没有太多帮助。我是从一个普遍的角度来谈论的。考虑到使用MP相当容易，尝试一下可能是值得的。与手动编写JIT的详细信息相比，尝试它的工作量要少得多。 - Richard