计算大量 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)

Answer 1

在评论中使用@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稍微快一些。

Answer 2

你得到了一些很好的建议，但由于这个特定的目标，我想再添加一个：

有没有办法显着加速这一过程？

实际上，如果您需要这些操作显着更快（通常意味着 > 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)

Answer 3

在一个回答中，尼克提到使用 GPU - 这当然是最好的解决方案。

但是 - 作为一般规则 - 您正在做的事情可能是 CPU 受限的。 因此（除了 GPU 方法），你能得到的最好结果是，如果你利用机器上的所有内核并行工作。

因此，为此您可能希望使用多处理（而不是 Python 的多线程！），将作业拆分为在每个内核上并行运行的多个部分。

这不是微不足道的，但也不是太难，网上有很多很好的例子/指南。

但是如果你有一台 8 核的机器，只要你小心地通过尝试在进程之间传递许多小对象来避免内存瓶颈，它可能会给你几乎 8 倍的速度增加，但是在一个组中将它们全部传递开始