没有获得高效的pytorch广播命令

Question

我的模型中有 5 classes 的 class-wise 特征向量。每个 class 的特征向量都是 20 维的。我想将标量增益乘以每个 class 的特征向量，并将得到的加权特征向量相加形成新的特征矩阵。

gain_matrix 对应于每个第 i j 对 classes 的标量值。第 i 个 class 的特征向量（20 维）计算为标量增益乘以所有其他 classes 特征向量的总和。具体实现代码如下。

nClass=5
feature_dim=20

gain_matrix=torch.rand(nClass,nClass)
feature_matrix=torch.rand(nClass,feature_dim) #in my implementation this is output from model
feature_matrix_new=torch.zeros(nClass,feature_dim)

for i in range(nClass):
    for j in range(nClass):
            feature_matrix_new[i,:]+=gain_matrix[i][j]*feature_matrix[j,:]

嵌套的 for 循环大大减慢了执行速度。

是否有任何有效的 PyTorch 广播解决方案来避免我的实现中的嵌套 for 循环？

我看过pytorch broadcasting web page，但对我帮助不大。

Answer 1

这将是一个使用的好地方torch.einsum:

>>> feature_matrix_new = torch.einsum('ij,jk->ik', gain_matrix, feature_matrix)

然而在这种情况下，这归结为 matrix multiplication:

>>> feature_matrix_new = gain_matrix @ feature_matrix

没有获得高效的pytorch广播命令

Efficient pytorch broadcasting command not obtained

python

numpy

pytorch

array-broadcasting