PyTorch中不同形状张量间高效计算两两L2范数距离的方法

你可以直接使用PyTorch内置的torch.cdist()函数实现单行高效计算，完全替代低效的for循环：

l2_dist = torch.cdist(b.to(torch.float32), a.to(torch.float32), p=2)

说明 #

• torch.cdist()专门用于计算两组张量元素间的两两距离，默认p=2对应L2范数，无需额外手动展开计算
• 输出张量形状为torch.Size([128, 1600])，和你循环实现的结果完全一致
• 该函数基于向量化操作和PyTorch底层优化，运行效率远高于手动for循环

如果想手动基于广播机制实现（原理与cdist一致），也可以用这行代码：

l2_dist = torch.sqrt(torch.sum((b.unsqueeze(1) - a.unsqueeze(0)) ** 2, dim=-1)).to(torch.float32)

内容的提问来源于stack exchange，提问作者Arun