使用 einsum 对 numpy 进行进一步优化以进行堆叠矩阵向量乘法

Question

我有一个相对简单的例子 "particle propagation" 使用线性矩阵变换。

我的粒子分布基本上是一组（'bunch'）5维向量。它通常包含 100k 到 1M 这样的向量。

这些向量中的每一个都必须乘以一个矩阵。

到目前为止我想出的解决方案如下。

粒子是这样创建的，协方差矩阵在这里显示为对角线，但这是为了一个相对简单的例子：

# Edit: I now use np.random_intel linking to MKL for improved performances
d = np.random.multivariate_normal(
    [0.0,
     0.0,
     0.0,
     0.0,
     0.0
     ],
    np.array([
        [1.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 0.0, 0.1]
    ]),
    int(1e5)
)

传播矩阵很简单

D = np.array([[1, 10, 0, 0], 
          [0, 1, 0, 0],
          [0, 0, 1, 0],
          [0, 0, 0, 1]])

我对einsum的解决方案是

r = np.einsum('ij,kj->ik', d[:, 0:4], D)

（注意这里我滑动只得到了向量的前四个坐标但是无关的原因）

Is there a way to make this significantly faster?

我不是很清楚所有的细节，但这里有一些想法：

einsum 默认情况下不调用 BLAS 但使用内部 SSE 优化，有没有一种方法可以用纯 BLAS 调用来表达我的问题，从而使它更快？
显然是 einsum 的 optimize 选项的最新版本，可以在更广泛的情况下打开以回退到 BLAS 调用。我试过了，它不会改变执行时间。
使用 PyPy 和 numpy 会更好吗？

我测试了@Divakar 的建议，它确实更快（10M 粒子）：

%%timeit
r = d[:, 0:4].dot(D.T)
# 541 ms ± 9.44 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

与我最初的相比

%%timeit -n 1 -r 1
r = np.einsum('ij,kj->ik', d[:, 0:4], D, optimize=True)
# 1.74 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

可能影响最终答案的直接相关问题：

How can I deal with 'lost' particles?

在逐个粒子矩阵乘法之后，我将检查一些坐标的上限，例如（r 是上一步的结果：

selected = (r[:, 0] < 0.1) & (r[:, 1] < 0.1)
ind = np.where(selected)
r[ind]

然后应用r[ind].

的下一轮矩阵乘法

有几件事我不太清楚：

这是最有效的吗？
它不会创建太多副本吗？
"keep" 未选择的粒子（并无论如何将它们相乘），同时跟踪它们丢失的事实（通过遮罩）不是更好吗？这是更多的乘法，但可以将所有内容保存在一个对象中，无需进一步分配并保持所有内容对齐？

Answer 1

为了进一步提高@Divakar 建议的代码的性能，我宁愿建议使用 PyTorch library. This would give you over 2 orders of magnitude speedup when compared to plain the dot product (np.dot()) using NumPy arrays（对于您的情况，从 ms 到微秒；稍后会详细介绍）

首先，我将演示如何在 NumPy 和 PyTorch. (Since PyTorch shares the same memory of NumPy ndarray 中执行此操作，我们不需要做额外的工作）

计时

# setup inputs
In [61]: d = np.random.multivariate_normal(
    ...:     [0.0,
    ...:      0.0,
    ...:      0.0,
    ...:      0.0,
    ...:      0.0
    ...:      ],
    ...:     np.array([
    ...:         [1.0, 0.0, 0.0, 0.0, 0.0],
    ...:         [0.0, 1.0, 0.0, 0.0, 0.0],
    ...:         [0.0, 0.0, 1.0, 0.0, 0.0],
    ...:         [0.0, 0.0, 0.0, 1.0, 0.0],
    ...:         [0.0, 0.0, 0.0, 0.0, 0.1]
    ...:     ]),
    ...:     int(1e5)
    ...: )

In [62]: d.dtype
Out[62]: dtype('float64')

In [63]: D = np.array([[1, 10, 0, 0], 
    ...:           [0, 1, 0, 0],
    ...:           [0, 0, 1, 0],
    ...:           [0, 0, 0, 1]], dtype=np.float64)
    ...:           

In [64]: DT = D.T

In [65]: DT.dtype
Out[65]: dtype('float64')


# create input tensors in PyTorch
In [66]: d_tensor = torch.DoubleTensor(d[:, 0:4])

In [67]: DT_tensor = torch.DoubleTensor(DT)

# float64 tensors
In [69]: type(d_tensor), type(DT_tensor)
Out[69]: (torch.DoubleTensor, torch.DoubleTensor)

# dot/matmul using `np.dot()`
In [73]: np_dot = np.dot(d[:, 0:4], DT)

# matmul using `torch.matmul()`
In [74]: torch_matmul = torch.matmul(d_tensor, DT_tensor)

# sanity check!! :)
In [75]: np.allclose(np_dot, torch_matmul)
Out[75]: True

现在是不同方法的时机！

In [5]: %timeit r = np.einsum('ij,kj->ik', d[:, 0:4], D)
2.63 ms ± 97.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [6]: %timeit r = d[:, 0:4].dot(D.T)
1.56 ms ± 47.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [7]: %timeit r = np.einsum('ij,kj->ik', d[:, 0:4], D, optimize=True)
2.73 ms ± 136 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# over 2 orders of magnitude faster :)
In [14]: %timeit torch_matmul = torch.matmul(d_tensor, DT_tensor)
87 µs ± 7.71 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

需要注意的一件重要事情是，我们需要在两个 NumPy ndarray and PyTorch Tensors. (Here I used np.float64 since np.random.multivariate_normal returned float64 values. So, I upcasted the D matrix to float64. Correspondingly, when creating PyTorch tensors, I used torch.DoubleTensor 中具有 相同的数据类型 ，这相当于 np.float64。这是一种数据类型匹配是必不可少的以获得相同的结果，特别是在处理浮点数时）。

因此，关键要点是 PyTorch Tensor operations are orders of magnitude faster than NumPy ndarray 操作。

使用 einsum 对 numpy 进行进一步优化以进行堆叠矩阵向量乘法

Further optimization with numpy using einsum for stacked matrix-vector multiplication

python

python-3.x

numpy

numpy-ndarray

dot-product