Pytorch快速计算余弦相似性矩阵

之前一直想找一个可以快速用矩阵相乘（可以直接GPU加速）计算余弦相似性矩阵的代码，总算找到了。代码是在参考代码的基础上进行了一些修改以适配自己的任务。

import torch


def calculate_cosine_similarity_matrix(h_emb, eps=1e-8):
    r'''

    '''
    # h_emb (N, M)
    # normalize
    a_n = h_emb.norm(dim=1).unsqueeze(1)
    a_norm = h_emb / torch.max(a_n, eps * torch.ones_like(a_n))

    # cosine similarity matrix
    sim_matrix = torch.einsum('bc,cd->bd', a_norm, a_norm.transpose(0,1))
    return sim_matrix


if __name__ == "__main__":
    x = torch.randn(100, 700)
    sim_matrix = calculate_cosine_similarity_matrix(x, eps=1e-8)
    # print(sim_matrix)
    y = torch.zeros((100,100))
    for i in range(100):
        for j in range(100):
            y[i,j] = torch.cosine_similarity(x[i].unsqueeze(dim=0), x[j].unsqueeze(dim=0))

    print(y-sim_matrix)

输出结果：

tensor([[ 0.0000e+00,  0.0000e+00, -1.8626e-09,  ..., -1.7462e-09,
         -2.2352e-08, -1.6764e-08],
        [ 0.0000e+00, -2.3842e-07, -3.7253e-09,  ..., -2.6077e-08,
         -9.3132e-09, -1.1642e-08],
        [-1.8626e-09, -3.7253e-09,  1.1325e-06,  ..., -5.3551e-09,
         -4.6566e-09, -1.8626e-08],
        ...,
        [-1.7462e-09, -2.6077e-08, -5.3551e-09,  ...,  5.9605e-07,
         -2.4214e-08, -2.2352e-08],
        [-2.2352e-08, -9.3132e-09, -4.6566e-09,  ..., -2.4214e-08,
          2.3842e-07,  9.3132e-10],
        [-1.6764e-08, -1.1642e-08, -1.8626e-08,  ..., -2.2352e-08,
          9.3132e-10,  3.5763e-07]])

可以说误差非常小了，非常之nice。
接下来讲讲为什么这么写，首先回顾一下余弦相似性（图源自百度百科）：

其实一开始我理解不了这个代码为什么可以实现（本人比较呆）。首先我们可以将输入看作是一个向量而不是一个矩阵，那么

a_n = h_emb.norm(dim=1)

得到的是该向量的二范数，接着用

a_norm = h_emb / torch.max(a_n, eps * torch.ones_like(a_n))

可以看作是$\frac{A}{||A||}$，至于为什么要加eps是防止0除。现在你再把原来的向量看回是一个矩阵，那么最后的

sim_matrix = torch.einsum('bc,cd->bd', a_norm, a_norm.transpose(0,1))

就好理解得到的是一个相似性矩阵了。