以PyTorch为单位计算批次成对伸角距离

我有两个张量，它们的形状都一样。我想计算成对的伸角距离

我尝试过的：

import torch
import geomloss  # pip install git+https://github.com/jeanfeydy/geomloss

a = torch.rand((8,4))
b = torch.rand((8,4))

geomloss.SamplesLoss('sinkhorn')(a,b)
# ^ input shape [batch, feature_dim]
# will return a scalar value

geomloss.SamplesLoss('sinkhorn')(a.unsqueeze(1),b.unsqueeze(1))  
# ^ input shape [batch, n_points, feature_dim]
# will return a tensor of size [batch] of distances between a[i] and b[i] for each i

然而，我想计算两两距离，其中合成张量的大小应为

[batch，batch]

。为了实现这一点，我尝试了以下方法来使用广播：

geomloss.SamplesLoss('sinkhorn')(a.unsqueeze(0), b.unsqueeze(1))

但我得到了这个错误信息：

ValueError:样本

x

和

y

应具有相同的批大小

---

因为文档中没有给出如何使用距离向前函数的示例。这里有一种方法，需要调用distance函数

batch

times

我们将逐行构造距离矩阵。行

i

对应于距离

a[i]b[0]

，

a[i]b[1]

，到

a[i]b[batch]

。为此，我们需要为每一行

i

，构造一个

（8x4）

张量

a[i]

的重复版本

这将有助于：

a_i = torch.stack(8*[a[i]], dim=0)

然后，我们计算

a[i]

和

b

中的每个批次之间的距离：

dist(a_i.unsqueeze(1), b.unsqueeze(1))

总共有

batch

行，我们可以构建最终的张量

堆栈


以下是完整的代码：

batch = a.shape[0]
dist = geomloss.SamplesLoss('sinkhorn')
distances = [dist(torch.stack(batch*[a[i]]).unsqueeze(1), b.unsqueeze(1)) for i in range(batch)]
D = torch.stack(distances)

因此，您希望在每个批次索引对
（i，j）
处有
a[i]
和
b[j]
之间距离的2D张量，对吗？@Ivan是的，非常感谢您的解决方案。我还考虑将基于循环的解决方案作为一种选择。谢谢你的帮助！。