python - 在 PyTorch 中找到一批向量之间的 jaccard 相似性

Question

我有一批形状为 (bs, m, n) 的向量(即维度为 mxn 的 bs 向量)。对于每个批处理，我想计算第一个向量与其余 (m-1) 个向量的 Jaccard 相似度

例子:

a = [
    [[3, 8, 6, 8, 7],
    [9, 7, 4, 8, 1],
    [7, 8, 8, 5, 7],
    [3, 9, 9, 4, 4]],

    [[7, 3, 8, 1, 7],
    [3, 0, 3, 4, 2],
    [9, 1, 6, 1, 6],
    [2, 7, 0, 6, 6]]
]

找出 a[:,0,:] 和 a[:,1:,:] 之间的成对 jaccard 相似度 即，

[3, 8, 6, 8, 7] with each of [[9, 7, 4, 8, 1], [7, 8, 8, 5, 7], [3, 9, 9, 4, 4]] (3 scores)
and
[7, 3, 8, 1, 7] with each of [[3, 0, 3, 4, 2], [9, 1, 6, 1, 6], [2, 7, 0, 6, 6]] (3 scores)

这是我试过的Jaccard函数

def js(la1, la2):
    combined = torch.cat((la1, la2))
    union, counts = combined.unique(return_counts=True)
    intersection = union[counts > 1]
    torch.numel(intersection) / torch.numel(union)

虽然这适用于大小不等的张量，但这种方法的问题是每个组合(张量对)中唯一值的数量可能不同，并且由于 PyTorch 不支持锯齿状张量，我无法处理一次批量处理向量。

如果我无法以预期的清晰度表达问题，请告诉我。在这方面的任何帮助将不胜感激

编辑:这是通过迭代第一维和第二维实现的流程。我希望有一个用于批处理的以下代码的矢量化版本

bs = 2
m = 4
n = 5
a = torch.randint(0, 10, (bs, m, n))
print(f"Array is: \n{a}")

for bs_idx in range(bs):
    first = a[bs_idx,0,:]
    for row in range(1, m):
        second = a[bs_idx,row,:]
        idx = js(first, second)
        print(f'comparing{first} and {second}: {idx}')

Answer 1

我不知道如何在 pytorch 中实现这一点，因为 AFAIK pytorch 不支持张量上的集合操作。在您的 js() 实现中，union 计算应该有效，但是 intersection = union[counts > 1] 没有给您正确的结果如果其中一个张量包含重复值。另一方面，Numpy 内置了对 union1d 和 intersect1d 的支持。您可以使用 numpy 向量化来计算成对的 jaccard 索引，而无需使用 for 循环:

import numpy as np

def num_intersection(vec1: np.ndarray, vec2: np.ndarray) -> int:
    return np.intersect1d(vec1, vec2, assume_unique=False).size

def num_union(vec1: np.ndarray, vec2: np.ndarray) -> int:
    return np.union1d(vec1, vec2).size

def jaccard1d(vec1: np.ndarray, vec2: np.ndarray) -> float:
    assert vec1.ndim == vec2.ndim == 1 and vec1.shape[0] == vec2.shape[0], 'vec1 and vec2 must be 1D arrays of equal length'
    return num_intersection(vec1, vec2) / num_union(vec1, vec2)
jaccard2d = np.vectorize(jaccard1d, signature='(m),(n)->()')

def jaccard(vecs1: np.ndarray, vecs2: np.ndarray) -> np.ndarray:
    """
    Return intersection-over-union (Jaccard index) between two sets of vectors.
    Both sets of vectors are expected to be flattened to 2D, where dim 0 is the batch
    dimension and dim 1 contains the flattened vectors of length V (jaccard index of 
    an n-dimensional vector and of its flattened 1D-vector is equal).
    Args:
        vecs1 (ndarray[N, V]): first set of vectors
        vecs2 (ndarray[M, V]): second set of vectors
    Returns:
        ndarray[N, M]: the NxM matrix containing the pairwise jaccard indices for every vector in vecs1 and vecs2
    """
    assert vecs1.ndim == vecs2.ndim == 2 and vecs1.shape[1] == vecs2.shape[1], 'vecs1 and vecs2 must be 2D arrays with equal length in axis 1'
    return jaccard2d(vecs1, vecs2)

这当然不是最优的，因为代码不在 GPU 上运行。如果我使用形状 (1, 10) 的 vecs1 和形状 的 vecs2 运行 jaccard 函数(10_000, 10) 我在我的机器上得到的平均循环时间为 200 ms ± 1.34 ms，这对于大多数用例来说应该足够快了。 pytorch 和 numpy 数组之间的转换非常便宜。

要将此函数应用于数组 a 的问题:

a = torch.tensor(a).numpy()  # just to demonstrate
ious = [jaccard(batch[:1, :], batch[1:, :]) for batch in a]
np.array(ious).squeeze()  # 2 batches with 3 scores each -> 2x3 matrix

# array([[0.28571429, 0.4       , 0.16666667],
#        [0.14285714, 0.16666667, 0.14285714]])

如果需要，对结果使用 torch.from_numpy() 再次获得 pytorch 张量。

---

更新:

如果你需要一个pytorch版本来计算Jaccard索引，我在torch中部分实现了numpy的intersect1d:

from torch import Tensor

def torch_intersect1d(t1: Tensor, t2: Tensor, assume_unique: bool = False) -> Tensor:
    if t1.ndim > 1:
        t1 = t1.flatten()
    if t2.ndim > 1:
        t2 = t2.flatten()
    if not assume_unique:
        t1 = t1.unique(sorted=True)
        t2 = t2.unique(sorted=True)
    # generate a m x n intersection matrix where m is numel(t1) and n is numel(t2)
    intersect = t1[(t1.view(-1, 1) == t2.view(1, -1)).any(dim=1)]
    if not assume_unique:
        intersect = intersect.sort().values
    return intersect

def torch_union1d(t1: Tensor, t2: Tensor) -> Tensor:
    return torch.cat((t1.flatten(), t2.flatten())).unique()

def torch_jaccard1d(t1: Tensor, t2: Tensor) -> float:
    return torch_intersect1d(t1, t2).numel() / torch_union1d(t1, t2).numel()

要向量化 torch_jaccard1d 函数，您可能需要查看 torch.vmap ，它允许您在任意批处理维度上对函数进行矢量化(类似于 numpy 的 vectorize)。 vmap 函数是一个原型(prototype)功能，在通常的 pytorch 发行版中尚不可用，但您可以使用每晚构建的 pytorch 来获取它。我还没有测试过，但这可能有效。