没有循环和内存错误的 Numpy 逐元素点积

Question

我正在用 numpy 处理一个简单的问题。我有两个矩阵列表——比如A,B——编码为形状分别为(n,p,q)和(n,q,r)的3D数组。

我想计算它们的逐元素点积，即 3D 数组 C 使得 C[i,j,l] = sum A[i,j,:] B[i,:,l]。这在数学上非常简单，但我必须遵循以下规则：

1) 我只能使用 numpy 函数（dot、tensordot、einsum 等）：没有循环和 cie。这是因为我希望它可以在我的 gpu（使用 cupy）上工作，并且循环在它上面很糟糕。我希望在当前设备上进行所有操作。

2) 由于我的数据可能非常大，通常 A 和 B 已经占用了几十 Mb 的内存，我不想构建任何形状比 (n,p,q),(n,q,r),(n,p,r) 更大的项目（没有中间 4D数组必须存储）。

例如，我找到的解决方案 there ，即使用：

C = np.sum(np.transpose(A,(0,2,1)).reshape(n,p,q,1)*B.reshape(n,q,1,r),-3)

在数学上是正确的，但意味着中间创建了一个 (n,p,q,r) 数组，这对我的目的来说太大了。

我也遇到过类似的问题

C = np.einsum('ipq,iqr->ipr',A,B)

我不知道底层的操作和构造是什么，但它总是会导致内存错误。

另一方面，有点天真，比如：

C = np.array([A[i].dot(B[i]) for i in range(n)])

在内存方面似乎还可以，但在我的 gpu 上效率不高：列表似乎是在 CPU 上构建的，并且将其重新分配给 gpu 很慢（如果有一种对 cupy 友好的方式来编写它，它将是一个不错的解决方案！）

感谢您的帮助！

Answer 1

你想要numpy.matmul (cupy version here)。 matmul 是一个“广播”矩阵乘法。

我认为人们已经知道 numpy.dot 语义很不稳定，并且需要广播矩阵乘法，但是在 python 获得 @ 运算符之前，引入更改的动力并不大。我看不到dot 有什么用处，但我怀疑更好的语义和使用A @ B 的便利性意味着dot 将随着人们发现新功能和运算符而失宠。

Answer 2

您试图避免的迭代方法可能不是那么糟糕。例如，考虑以下时间：

In [51]: A = np.ones((100,10,10))
In [52]: timeit np.array([A[i].dot(A[i]) for i in range(A.shape[0])])
439 µs ± 1.35 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [53]: timeit np.einsum('ipq,iqr->ipr',A,A)
428 µs ± 170 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [54]: timeit A@A
426 µs ± 54.6 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

在这种情况下，所有三个都需要大约相同的时间。

但是我把后面的维度加倍，迭代的方式其实更快：

In [55]: A = np.ones((100,20,20))
In [56]: timeit np.array([A[i].dot(A[i]) for i in range(A.shape[0])])
702 µs ± 1.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [57]: timeit np.einsum('ipq,iqr->ipr',A,A)
1.89 ms ± 1.63 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [58]: timeit A@A
1.89 ms ± 490 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

当我将 20 更改为 30 和 40 时，同样的模式成立。我有点惊讶 matmul 时间与 einsum 如此接近。

我想我可以尝试将这些推到内存限制。我没有花哨的后端来测试这方面。

一旦考虑到内存管理问题，对一个大问题进行少量迭代并不是那么可怕。在 numpy 中，您想要避免的事情是对一个简单任务进行多次迭代。