批量相似度计算

最新推荐文章于 2022-11-14 20:57:58 发布

curry3030

最新推荐文章于 2022-11-14 20:57:58 发布

阅读量570

点赞数

CC 4.0 BY-SA版权

文章标签：深度学习 pytorch 人工智能

本文链接：https://blog.youkuaiyun.com/curry3030/article/details/125024298

本文介绍了一种利用PyTorch库实现的批量余弦相似度计算方法，该方法能够高效地计算不同尺寸向量之间的相似度，避免了使用循环所带来的效率问题。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import torch

def batch_cosine_similarity(sentence_embedding_a, sentence_embedding_b):
    """
    :param sentence_embedding_a: [a_number, hidden_dim]
    :param sentence_embedding_b: [b_number, hidden_dim]
    :return:
    """
    a_number = sentence_embedding_a.size(0)
    b_number = sentence_embedding_b.size(0)

    # a_embedding: [a_number, hidden_dim] --> [a_number, b_number, hidden_dim] --> [a_number * b_number, hidden_dim]
    # b_embedding: [b_number, hidden_dim] --> [a_number, b_number, hidden_dim] --> [a_number * b_number, hidden_dim]
    a_embedding = sentence_embedding_a.unsqueeze(1).repeat(1, b_number, 1).view(-1, 768)
    b_embedding = sentence_embedding_b.unsqueeze(0).repeat(a_number, 1, 1).view(-1, 768)

    # similarity: [b_number * a_number, 1]
    similarity = torch.cosine_similarity(a_embedding, b_embedding)
    similarity = similarity.view(a_number, b_number)

    return similarity

torch 自带的 cosine_similarity 函数似乎仅适用于相同尺寸的向量进行相似度的计算，为了避免for循环带来的速度慢的情况，采用repeat方式进行向量的复制，并且实现多对多的的余弦相似度计算。