Python 线性空间中的Numpy矩阵乘法极大值

我有两个大数组

A

is

nxk

和

B

is

kxm

我想计算

np.argmax（A@B，axis=1）

，这是每行的最大列

不幸的是，如果我天真地这样做，numpy将计算整个数组

A@B

，这需要

nxm

内存——比我的内存还多

不需要额外的内存就可以做到这一点，只需单独计算每个条目并保持最大值

---

在numpy中可以做到这一点吗？

正如评论中所讨论的，numpy不能完全做到您想要的。但是，如果您有足够的内存，您可以迭代

A

的行，并逐段执行矩阵乘法，在循环的每次迭代中为生成的

1xk

收集argmax。这将是完全使用numpy和从头开始滚动您自己的实现之间的一种折衷。类似下面的内容。你也可以用麻木来提速

import numpy as np
from numba import jit

a = np.random.randn(500,200)
b = np.random.randn(200,1000)

def max_by_row(a,b):
    out = np.zeros((a.shape[0],),dtype='int64')
    for idx in range(a.shape[0]):
        out[idx] = np.argmax(a[idx,:] @ b)
    return out
        
@jit        
def max_by_row_jit(a,b):
    out = np.zeros((a.shape[0],),dtype='int64')
    for idx in range(a.shape[0]):
        out[idx] = np.argmax(a[idx,:] @ b)
    return out  

比较时间：

%timeit np.argmax(a@b,axis=1)
4.6 ms ± 226 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit max_by_row(a,b)
12.2 ms ± 233 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit max_by_row_jit(a,b)
8.85 ms ± 135 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

最后我用了这个。它采用两个矩阵，

A

和

B

，对于

A

的每一行

A

，它告诉您

B

的哪一行更近

---

这为我节省了20倍的时间，甚至不考虑内存，在我的例子中，

B

大约有一百万行。

像“在numpy中可以这样做吗？”这样的问题并不好。是的，这是可能的，你已经描述了如何，那么是什么阻止了你？我想你是在问numpy是否实现了任何惰性矩阵乘法，我认为它没有。你必须自己写出逻辑。还有，你是指线性空间吗？我不知道任何线性时间矩阵乘法算法。@juanpa.arrivillaga我错了，我自然是指线性空间。阻止我的是我描述的方法的二次空间要求。我可以用python循环编写逻辑，但这就消除了使用numpy.Right的性能优势。不幸的是，我不认为

numpy

实现了您想要的（不是100%确定）。但是，您可以尝试使用

numba

，一种用于基于numpy代码的JIT编译器。注意，naive算法在算法上比引擎盖下可能使用的

numpy

要慢，而且实际上，可能会根据BLAS/LAPACK后端的质量以各种神秘的方式进行优化。但这种幼稚但懒惰的方法可能适合你。幸运的是，

numba

非常擅长加速基于循环的代码，所以如果我是你，我会尝试一下。谢谢。纳巴很有趣。不过，我认为这一次对我来说还不够，但我会记住这一点。