Pytorch 计算一个矩阵行向量之间的相似度

import torch


def similarity(input_tensor):
    # step 1. 计算行向量的长度
    len_a = torch.sqrt(torch.sum(input_tensor ** 2, dim=-1))
    b = len_a.unsqueeze(1).expand(-1, input_tensor.shape[0])
    c = len_a.expand(input_tensor.shape[0], -1)
    # step2. 计算乘积
    x = input_tensor @ input_tensor.T
    # step3. 计算最后的结果
    res = x / (b * c)
    return res


def torch_similarity(input_tensor):
    sim_bg=torch.cosine_similarity(input_tensor.detach().unsqueeze(0), input_tensor.detach().unsqueeze(1),dim=-1)
    return sim_bg

if __name__ == '__main__':
    x=torch.randn(4,512)
    sim_bg=torch_similarity(x)
    print(sim_bg)
    res=similarity(x)
    print(res)
    print(res.size())

方法一：运算速度快