深度学习pytorch计算cosine距离

先对数据进行L2归一化再计算点积就是余弦距离

pytorch

import torch.nn.functional as F
def cosine_similarity_matrix(A,B):
    """
    计算矩阵 B 中每一行和矩阵 A 中每一行的余弦相似度
    参数:
        A: 矩阵 A, 大小为 (m, n)
        B: 矩阵 B, 大小为 (k, n)
    返回:
        cos_sim_matrix: 余弦相似度矩阵, 大小为 (k, m)
    """
    # 矩阵A,B进行L2归一化
    A_normlized = F.normlize(A,p=2,dim=1)
    B_normlized = F.normlize(B,p=2,dim=1)
    cos_sim_matrix = torch.matul(B_normlized,A_normlized.T)
    return cos_sim_matrix

numpy

import numpy as np

def cosine_similarity_matrix(A, B):
    """
    计算矩阵 B 中每一行和矩阵 A 中每一行的余弦相似度
    参数:
        A: 矩阵 A, 大小为 (m, n)
        B: 矩阵 B, 大小为 (k, n)
    返回:
        cos_sim_matrix: 余弦相似度矩阵, 大小为 (k, m)
    """
    # 矩阵A,B进行L2归一化
    A_normalized = A / np.linalg.norm(A, axis=1, keepdims=True)
    B_normalized = B / np.linalg.norm(B, axis=1, keepdims=True)
    # 计算余弦相似度
    cos_sim_matrix = np.dot(B_normalized, A_normalized.T)
    return cos_sim_matrix