向量化的成对距离答案

【问题标题】：Vectorised pairwise distance向量化的成对距离
【发布时间】：2022-06-16 18:53:23
【问题描述】：

TLDR：给定两个张量t1 和t2，它们代表形状为c,h,w 的张量的b 样本（即，每个张量的形状为b,c,h,w），我'我试图有效地计算所有i,j 的t1[i] 和t2[j] 之间的成对距离

更多上下文 - 我已经为我的训练和测试数据 (CIFAR10) 提取了 ResNet18 激活，并且我正在尝试实现 k-nearest-neighbours。一个可能的伪代码可能是：

for te in test_activations:
    distances = []
    for tr in train_activations:
        distances.append(||te-tr||)
    neighbors = k_smallest_elements(distances)
    prediction(te) = majority_vote(labels(neighbors))

我正在尝试对来自 test 和 train 激活数据集的批次进行矢量化处理。我试过迭代批次（而不是样本）并使用torch.cdist(train_batch,test_batch)，但我不太确定这个函数如何处理多维张量，如documentation 中所述

torch.cdist(x1, x2,...):
如果x1 的形状为BxPxM，x2 的形状为BxRxM，则输出的形状为BxPxR

这似乎无法处理我的情况（见下文）

可以在此处找到一个最小示例：

b,c,h,w = 1000,128,28,28 # actual dimensions in my problem
train_batch = torch.randn(b,c,h,w)
test_batch = torch.randn(b,c,h,w)
d = torch.cdist(train_batch,test_batch)

您可以将test_batch 和train_batch 视为for 循环中的张量for test_batch in train: for train_batch in test:...

【问题讨论】：

标签： python machine-learning pytorch knn nearest-neighbor

【解决方案1】：

在将数据提供给内置 PyTorch 运算符之前，通常必须对其进行整形。正如您所说，torch.cdist 使用两个形状为(B, P, M) 和(B, R, M) 的输入，并返回一个形状为(B, P, R) 的张量。

相反，您有两个形状相同的张量：(b, c, h, w)。如果我们匹配这些维度：B=b、M=c，而P=h*w（来自第一个张量）和R=h*w（来自第二个张量） .这需要将空间维度展平在一起并交换最后两个轴。比如：

>>> x1 = train_batch.flatten(2).transpose(1,2)
>>> x2 = test_batch.flatten(2).transpose(1,2)

>>> d = torch.cdist(x1, x2)

现在d 包含所有可能对之间的距离(train_batch[b, :, iy, ix], test_batch[b, :, jy, jx]) 并且形状为(b, h*w, h*w)。

然后，您可以使用 argmax 应用 knn 来检索从训练批次的一个元素到测试批次的 k 最近邻。

【讨论】：