PyTorch 中一组向量之间的成对相似度矩阵

Question

假设我们有一个 3D PyTorch 张量，其中第一个维度表示batch_size，如下所示：

import torch
import torch.nn as nn
x = torch.randn(32, 100, 25)

也就是说，对于每个i，x[i]是一组 100 个 25 维向量。我想计算每个批次项目的这些向量的相似度（例如，余弦相似度——但一般来说是任何这样的成对距离/相似度矩阵）。

也就是说，对于每个x[i]我需要计算一个[100, 100]矩阵，其中包含上述向量的成对相似度。x[t]更具体地说，对于所有，该矩阵的第 (i,j) 个元素应包含 (100x25) 的第 i 行和第 j 行之间的相似度（或距离）t=1, ..., batch_size。

如果我使用torch.nn.CosineSimilarity()，无论dim我使用什么，结果都是[100, 25]( dim=0)，或[32, 25]( dim=1) ，其中我需要大小为 的张量[32, 100, 100]。我希望torch.nn.CosineSimilarity()以这种方式工作（因为，至少对我来说，它看起来更直观） ，但事实并非如此。

可以使用类似下面的东西来完成吗？

torch.matmul(x, x.permute(0, 2, 1))

我猜这可以给出距离矩阵，但是如果我需要任意成对运算怎么办？我应该使用上面的内容来构建这个操作吗？

或者也许我应该x以某种方式重复，以便我可以使用内置的torch.nn.CosineSimilarity()？

谢谢。

Answer 1

文档暗示输入的形状cosine_similarity必须相等，但事实并非如此。PyTorch 内部通过 进行广播torch.mul，插入带有切片（或torch.unsqueeze）的维度将为您提供所需的结果。由于上下三角形的重复计算和内存，这不是最佳的，但它很简单：

import torch
from torch.nn import functional as F
from scipy.spatial import distance

# compute once in pytorch
x = torch.randn(32, 100, 25)
y = F.cosine_similarity(x[..., None, :, :], x[..., :, None, :], dim=-1)

assert y.shape == torch.Size([32, 100, 100])

# test against scipy by iterating over each batch element
z = []
for i in range(x.shape[0]):
    slice = x[i, ...].numpy()
    z.append(torch.tensor(distance.cdist(slice, slice, metric='cosine'), dtype=torch.float32))

# convert similarity to distance and ensure they're reasonably close
assert torch.allclose(torch.stack(z), 1.0-y)