场景
给定一个三维矩阵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)