高效的矩阵-矩阵乘积答案

【问题标题】：Efficient matrix-matrix product高效的矩阵-矩阵乘积
【发布时间】：2020-05-22 10:11:23
【问题描述】：

我正在尝试在多级贝叶斯模型中计算矩阵矩阵乘积。数据和模型变量的形状如下：

X_train: (n_samples, n_predictors)
alpha: (n_groups, )
beta: (n_predictors, n_groups)

X_train 是密集的。 alphas 是组截距，betas 表示斜率（同样，每个组的变量）。样本被分成几组；数组group_index（大小为(n_samples,)）表示每个样本所属的组。简而言之，线性模型是

y[n] = alpha[group_index[n]] + < X_train[n, 1:K], beta[1:K, group_index[n]] >

对于所有n = 1 ... n_samples，其中< , > 表示内积。

以下是在 Python 中的实现方式：

# an element-wise product between two matrices of size [n_samples, n_predictors]
y = alpha[group_index] + (X_train * beta[:, group_index].T).sum(axis=-1)

问题：这个实现可以提高内存/cpu 的效率吗？

这里n_samples 可以运行到100 万，而n_predictors 和n_groups 大约是100 左右。我正在寻找的是一个矢量化的公式，它 (1) 需要更少的存储空间，并且 (2) 与我上面提出的相比，它的运行速度差不多。

【问题讨论】：

评论说“两个大小为 [n_samples, n_predictors] 的矩阵之间的元素乘积”，但这似乎不是代码中发生的事情。
@user2357112supportsMonica 我认为是发生了什么：X_samples 是 (n_samples, n_predictors)，beta 是 (n_predictors, n_groups) 和 group_index 是 (n_samples,)。因此beta[:, group_index].T 的大小为(n_samples, n_predictors)。
啊，group_index 是一个数组。

标签： python numpy machine-learning vectorization linear-regression

【解决方案1】：

如果您将 X_train 重塑为 (N_samples,1,N_predictors) 并将 beta 重塑为 (1,N_predictors,N_groups)，那么您可以计算它们之间的矩阵乘法，得到 (N_samples,1,1)，您可以在添加到 alpha[group_index] 之前将其展平。

X_train_r = np.expand_dims(X_train,1)
beta_r = np.expand_dims(beta,0)
y = alpha[group_index] + (X_train_r @ beta_r[:,:, group_index].T).ravel()

它的计算速度也应该快近两倍。

【讨论】：

我并不清楚这个解决方案如何减少我的内存需求。 beta_r[:,:, group_index].T 的形状为 (N_samples, 1, N_predictors)
另外，这个解决方案似乎稍微快一些（在 %timeit 测试中大约为 5%）
是的，它并没有完全减少存储空间，因为它仍然使用来自 X_train 和 beta 的所有值。效率增益在矩阵乘法中；不是将两个形状为 (1e6,100) 的矩阵相乘，然后对 100 列（每列长 1e6 个值）进行列求和，而是 (1e6,1,100) 和 (1e6,100,1) 之间的 matmul 使用较少的乘法累加运算. 1e6 nsample 和 100 的 %timeit 测试分别为 859 和 468 毫秒。

【解决方案2】：

要尝试的另一件事是：

alpha[group_index] + np.einsum('ij,ji->i', X_train_r, beta[:, group_index])

用一个小例子进行测试：

In [192]: A = np.arange(10)                                                              
In [193]: B = np.arange(50).reshape(5,10)                                                
In [194]: X = np.arange(50).reshape(10,5)                                                

In [195]: A + (X * B.T).sum(axis=-1)                                                     
Out[195]: array([ 300,  836, 1422, 2058, 2744, 3480, 4266, 5102, 5988, 6924])
In [196]: A + np.einsum('ij,ji->i',X,B)                                                  
Out[196]: array([ 300,  836, 1422, 2058, 2744, 3480, 4266, 5102, 5988, 6924])

和时间：

In [197]: timeit A + (X * B.T).sum(axis=-1)                                              
14.1 µs ± 68 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [198]: timeit A + np.einsum('ij,ji->i',X,B)                                           
9.66 µs ± 9.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

和 Mercury 的解决方案：

In [208]: A + (X[:,None,:]@B.T[:,:,None]).ravel()                                        
Out[208]: array([ 300,  836, 1422, 2058, 2744, 3480, 4266, 5102, 5988, 6924])
In [209]: timeit A + (X[:,None,:]@B.T[:,:,None]).ravel()                                 
7.3 µs ± 18.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

matmul 针对矩阵乘积的批处理进行了优化，将数组传递给快速数学库（BLAS 等）。

【讨论】：