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。

接下來講講為什么這么寫

首先回顧一下余弦相似性（圖源自百度百科）：

其實一開始我理解不了這個代碼為什么可以實現(xiàn)（本人比較呆）。

首先我們可以將輸入看作是一個向量而不是一個矩陣，那么

a_n = h_emb.norm(dim=1)

得到的是該向量的二范數(shù)，接著用

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

可以看作是，至于為什么要加eps是防止0除。

現(xiàn)在你再把原來的向量看回是一個矩陣，那么最后的

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

就好理解得到的是一個相似性矩陣了。

總結

以上為個人經驗，希望能給大家一個參考，也希望大家多多支持腳本之家。

最新評論