pytorch 或 numpy 中的稀疏和巨大矩阵乘法

Question

我在大型稀疏矩阵上执行矩阵乘法时面临性能问题。我有一个场景，我需要将一个小尺寸向量

a

与一个巨大且高度稀疏的向量

b

相乘。这是代码的简化版本：

import numpy as np

B = 32
M = 10000000

a = np.random.rand(B)
b = np.random.rand(B, M)
b = b > 0.9

result = a @ b

在我的实际用例中，由于尺寸较大，

b

向量是从

np.memmap

文件加载的。

上述矩阵乘法运算需要大量时间来计算。我正在寻求有关如何优化此矩阵乘法运算以提高其性能的建议。任何见解或代码示例将不胜感激。

Answer 1

我认为你可以使用

np.einsum

。

import numpy as np
import time

B = 32
M = 10000000

a = np.random.rand(B)
b = np.random.rand(B, M)
b = b > 0.9

start = time.time()
result = a @ b
end = time.time() - start
print(end)

start = time.time()
a = np.expand_dims(a, 0)
result = np.einsum('ab,bm->am', a, b)
end = time.time() - start
print(end)
print(np.allclose(result, result))

结果

0.7853567600250244
0.4943866729736328
True