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等）：no loop＆amp; CIE。这是因为我希望这可以在我的gpu上工作（带有杯状），并且循环很糟糕。我希望在当前设备上进行所有操作。

2）由于我的数据可能很大，通常A和B已经在内存中占用了几十Mb，我不想构建形状比{更大}的任何项目{1}}（不能存储中间4D数组）。

例如，我找到的解决方案是there，正在使用：

(n,p,q),(n,q,r),(n,p,r)

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

我遇到过像

这样的问题

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

我不知道什么是潜在的操作＆amp;结构，但它总是导致内存错误。

另一方面，有些天真的事情如下：

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

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

感谢您的帮助！

Answer 1

您想要numpy.matmul（cupy version here）。 matmul是一个“广播”矩阵乘法。

我认为人们已经知道numpy.dot语义是不可靠的，并且需要广播矩阵乘法，但是在python获得@运算符之前，引入变化的动力并不大。我没有看到dot在任何地方，但我怀疑更好的语义和执行A @ B的容易性将意味着当dot发现新的函数和运算符时{{1}}将失宠

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 element-wise dot产品没有循环和内存错误

2 个答案: