使用einsum进行堆积矩阵向量乘法的numpy进一步优化

Question

我有一个相对简单的使用线性矩阵变换的“粒子传播”的情况。

我的粒子分布基本上是一组（'束'）的五维向量。它通常包含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)

（注意，我在这里滑动以获得向量的前四个坐标，但不相关的原因）。

有没有办法让这个显着加快？

我对所有细节都没有清晰的看法，但这里有一些想法：

• 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)

---

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

我该如何处理“丢失”的粒​​子？

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

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

然后应用r[ind]的下一轮矩阵乘法。

有些事情我不清楚：

• 这是最有效的吗？
• 它不会创造太多副本吗？
• 最好“保留”未选择的粒子（无论如何都要乘以它们），同时跟踪它们丢失的事实（通过掩码）？这是一个更多的乘法，但这可以将所有内容保存在一个对象中，没有进一步的分配并保持一切对齐？

Answer 1

为了进一步提高@Divakar建议的代码的性能，我宁愿建议使用PyTorch库。与普通的点积（np.dot()）相比，使用NumPy arrays（对于你的情况，从ms到微秒;相关的更多内容），这将为你提供超过2个数量级的加速

首先，我将演示如何在NumPy和PyTorch中完成它。 （因为PyTorch和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和PyTorch Tensors中使用相同的数据类型。 （这里我使用np.float64，因为np.random.multivariate_normal返回float64值。所以，我将D矩阵升级到float64。相应地，当创建PyTorch张量时，我使用torch.DoubleTensor相当于np.float64。这种数据类型匹配对于获得相同的结果是必不可少的，特别是在处理浮点数时）。

---

因此，关键的一点是，PyTorch Tensor的运营速度比NumPy ndarray运营快几个数量级。