如何使用行-时间-列对近似矩阵乘法进行矢量化

Question

我有两个矩阵A和B，我想使用行和列的子集来计算它们的AMM(近似矩阵乘积(。假设对于A的行和B的列，我分别具有两个概率分布pr和pc。A和B的尺寸均为900 x 900

我从具有概率p的行分布pr的A中选择一行，并且从具有概率的分布pc的B中选择一列(1-p(。

这意味着我从A中选择kp行，从B中选择k(1-p(列，其中k < D

下面给出了一个示例代码，

import numpy as np
D = 1000
p = 1/2
k = 500
# Assume uniform distribution for this case
pr = np.full(D, 1/D)
pc = pr
indexR = np.random.choice(1000, int(k*p)+1, p=pr)
indexC = np.random.choice(1000, int(k*(1-p))+1, p=pc)
A = np.random.uniform(0,1,(D,D))
B = np.random.uniform(0,10,(D,D))
# True multiplication
C = A @ B
C_hat = np.zeros((D,D))
# Approx. Multiplication
for r in indexR:
for c in indexC:
C_hat[r,c] = np.inner(A[r,:],B[:,c])

如何删除这些循环？

Answer 1

indexR和indexC可以包含重复项。因此，最好先把它们去掉排序索引也有助于提高性能，因为现代处理器确实不喜欢随机内存访问(因为它们不能从RAM预取数据，因为地址不能预测(。那么np.inner(A[r,:],B[:,c])只不过是一个点积，您可以将许多运算的点积重写为简单的向量矩阵积或矩阵乘法。

unique_indexR = np.unique(indexR)
unique_indexC = np.unique(indexC)
c_indices = np.tile(unique_indexC, len(unique_indexR))
r_indices = np.repeat(unique_indexR, len(unique_indexC))
C_hat = np.zeros((D,D))
C_hat[r_indices, c_indices] = (A[unique_indexR,:] @ B[:,unique_indexC]).reshape(-1)

以下是性能结果(在我的6核机器上(：

Initial time:         261.46 ms
Vectorized version:     3.67 ms

因此，矢量化代码的速度是的71倍。