大矩阵的余弦相似度

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. pairwise_distance 完全相同。