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的{{1}}完全相等。

大矩阵的余弦相似度

1 个答案: