pytorch 余弦相似度矩阵cos_similar，批量运算_批量计算余弦相似度 python

import torch
from torch import Tensor
 
 
# params: 张量p和q，保证最后一个维度为特征维度，倒数第二哥维度为所求的相似维度。 
# return: sim_matrix[i][j]代表p的第i个特征和q的第j个特征相似度计算值。
def cos_similar(p: Tensor, q: Tensor):
    sim_matrix = p.matmul(q.transpose(-2, -1))
    a = torch.norm(p, p=2, dim=-1)
    b = torch.norm(q, p=2, dim=-1)
    sim_matrix /= a.unsqueeze(-1)
    sim_matrix /= b.unsqueeze(-2)
    return sim_matrix
 
 
if __name__ == '__main__':
    a = torch.rand((64, 23, 50, 200))
    b = torch.rand((64, 23, 60, 200))
    print(cos_similar(a, b).size())