了解不同维度张量的广播和算术运算

Question

我目前正在计算向量之间的各种相似性指标，例如余弦相似性、欧氏距离、马哈拉诺比斯距离等。由于我处理的向量可能非常大，因此我需要将计算时间降至最低.

我正在努力了解如何使用不同维度的向量（但是，它们共享一个维度）以及如何在 PyTorch 中使用它。

我有两个向量，A 和 B，尺寸分别为 [867, 768] 和 [621, 768]。

我正在尝试计算以下内容：

v_a

v_a - v_b

我知道这可以通过 scipy 和 numpy 等实现，但我试图避免分离向量并将其移动到 CPU 以提高速度。

谁能帮我理解 PyTorch 中实现此目的所需的运算符逻辑？

Answer 1

您可以在两个输入张量上使用花式索引来取消压缩维度，使得 A 和 B 分别具有 (1, 867, 768) 和 (621, 1, 768) 的形状。减法运算将自动将两个张量传播到相同的形状。

>>> diff = A[None]-B[:,None]
>>> diff.shape
torch.Size([621, 867, 768])

这是实现 batched-pairwise 距离的典型方法。

更具体地说，请注意两个变体之间的区别：A[None]-B[:,None] 和 A[:,None]-B[None]。

diff = A[None]-B[:,None]    # (1, k_a, b) - (k_b, 1, b) -> (k_b, k_a, b) - (k_b, k_a, b)
diff.shape                  # (k_b, k_a, b)

相比于：

diff = A[:,None]-B[None]    # (k_a, 1, b) - (1, k_b, b) -> (k_a, k_b, b) - (k_a, k_b, b)
diff.shape                  # (k_a, k_b, b)

您可以在 NumPy documentation page 上阅读有关广播的更多信息。

Understanding broadcasting and arithmetic operations on different dimension tensors