torch三维矩阵中求最后一个维度所有向量两两之间的余弦相似度

场景

给定一个三维矩阵x(batch, seq_len, input_size)，最后需要得到一个余弦相似度矩阵e(batch_size, seq_len, seq_len)，例如e[0, 1, 2]=cos(x[0, 1, :], x[0, 2, :])。

实现

x = torch.rand(64, 24, 7)
e = torch.cosine_similarity(x.unsqueeze(2),
                            x.unsqueeze(1),
                            dim=-1)

假如令a=x.unsqueeze(2)=(64, 24, 1, 7)，b=x.unsqueeze(1)=(64, 1, 24, 7)，dim=-1表示在最后一维进行余弦计算，即a的每个(1, 7)都将与b中的(24, 7)计算余弦，最终计算 24 × 24 24\times24 24×24次。验证：

print(cos.shape)
print(F.cosine_similarity(x[0, 7, :], x[0, 8, :], dim=0))
print(e[0, 7, 8])
torch.Size([64, 24, 24])
tensor(0.5933)
tensor(0.5933)