未获得高效的pytorch广播命令

Question

我的模型中有 5 类的分类特征向量。每个类的特征向量是20 维。我想将标量增益乘以每个类的特征向量，然后将得到的加权特征向量相加以形成新的特征矩阵。

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

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