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

Question

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

我的粒子分布基本上是一组（“束”）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]])

我对@9​​87654323@ 的解决方案是

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 库。与使用 NumPy arrays 的普通 点积 (np.dot()) 相比，这将为您提供超过 2 个数量级的加速（对于您的情况，从 ms 到微秒；稍后会详细介绍）

首先，我将演示如何在 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 张量中具有相同的数据类型。 （这里我使用np.float64，因为np.random.multivariate_normal返回float64值。所以，我将D矩阵向上转换为float64。相应地，在创建PyTorch张量时，我使用了torch.DoubleTensor，它相当于@ 987654342@。这是一种数据类型匹配是必不可少的，以获得相同的结果，尤其是在处理浮点数时）。

---

因此，关键点是PyTorch Tensor 操作比NumPy ndarray 操作快几个数量级。