Pytorch 中的torch.einsum 求cos 相似度矩阵

einsum 计算consine similarity for contrastive learning

consine similarity 定义
$$S_{ij}=X_i^T{Y}_j/\Vert X_i{\Vert }_2\Vert Y_j{\Vert }_2$$
X 和Y为embedding 矩阵:$bz\times hiden\_dim$

$$X_i^T{Y}_j=\sum _k{X}_{ik}{Y}_{jk}$$

这一步用爱因斯坦求和便是

 torch.einsum('ik,jk->ij', X, Y)

$$N_{ij}=\Vert X_i{\Vert }_2\Vert Y_j{\Vert }_2$$

 torch.einsum('i,j->ij', X.norm(dim=-1), Y.norm(dim=-1))

3. 最后计算实例

对比学习里边最常使用的x, 以及x_argumentation 计算相似度矩阵

x_abs = x.norm(dim=1)
x_aug_abs = x_aug.norm(dim=1)

sim_matrix = torch.einsum('ik,jk->ij', x, x_aug) / torch.einsum('i,j->ij', x_abs, x_aug_abs)

其他用法

# 矩阵按列求和
torch_ein_out = torch.einsum('ij->j', [a]).numpy()

https://zhuanlan.zhihu.com/p/361209187