python - 大矩阵上的余弦相似度

Question

我正在尝试计算大型矩阵(3m x 2048)的所有对之间的余弦距离，并使用 pytorch 提取前 30 个相似向量。以下是我的代码，运行良好，但每次迭代大约需要 30 秒，这对于 300 万个单词向量来说太长了。有什么办法可以加快速度吗？

import torch.nn.functional as F
import torch
from tqdm import tqdm
import gc



sym_dict={}
tmp_list=[]

tot_dict=torch.load('xbx.pt')


all_tensors = torch.cat([v.unsqueeze(0) for k,v in tot_dict.items()], dim=0)
token_list= [i for i in tot_dict.keys()]

del tot_dict
gc.collect()



for counter ,value in tqdm(enumerate(token_list)):


  uniq_vec=torch.unsqueeze(all_tensors[counter],dim=0)

  dist = 1 - F.cosine_similarity(uniq_vec,all_tensors)
  index_sorted = torch.argsort(dist)


  roll_me=index_sorted[:30].cpu().numpy().tolist()

  for ind in roll_me:
    tmp_list.append(token_list[ind])
  sym_dict.update({value:tmp_list})
  tmp_list=[]



#save .pt file
torch.save(sym_dict,'sym_dict.pt')

Answer 1

直接求两个矩阵之间的成对距离可以吗？代码如下:

def pairwise_dist(x, y,p=2, eps=1e-6):
    x_a  =x[..., None, :, :]
    y_a  =y[...,None,:]
    dist = torch.pow(torch.abs((x_a - y_a) + eps), p).sum(dim=-1, keepdim=True).squeeze(2)
    return torch.pow(dist, 1/p)

t1 = torch.rand(3, 10)
t2 = torch.rand(4,10)
dist = pairwise_dist(t1,t2, eps=0)
print(dist)

dist 的形状为 4 x 3，其中每行表示 t1 的所有向量与 t2< 向量的距离.

请注意，这里两个向量之间的成对距离与 Pytorch 的 F 完全相同。成对距离。