查找PyTorch中一批向量之间的jaccard相似性



我有一批形状为(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:,:]之间的成对纸牌相似度例如,

[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}')

我不知道如何在pytorch中实现这一点,因为AFAIK pytorch不支持张量的集合操作。在您的js()实现中,union计算应该工作,但如果其中一个张量包含重复值,intersection = union[counts > 1]不会给您正确的结果。另一方面,Numpy内置了对union1dintersect1d的支持。您可以使用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和形状为(10_000, 10)vecs2运行jaccard函数,我在我的机器上得到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函数是一个原型特性,在通常的pytorch发行版中尚未提供,但您可以使用pytorch的夜间构建来获得它。我还没有测试过,但这可能有效。

您可以使用numpy。独特,愚蠢,交叉,愚蠢。Union1d和transform推荐使用python计算jaccard_similarity的函数:

import torch
import numpy as np
def jaccard_similarity_numpy(list1, list2):
s1 = np.unique(list1)
s2 = np.unique(list2)
return (len(np.intersect1d(s1,s2))) / len(np.union1d(s1,s2))
tns_a = torch.tensor([
[[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]]
])
tns_b = torch.empty(tns_a.shape[0],tns_a.shape[1]-1)
for i in range(tns_a.shape[0]):
tns_tmp = tns_a[i][0]
for j , tns in enumerate(tns_a[i][1:]):
tns_b[i][j] = (jaccard_similarity_numpy(tns_tmp, tns))
print(tns_b)

输出:

tensor([[0.2857, 0.4000, 0.1667],
[0.1429, 0.1667, 0.1429]])

你不标记numba,但是你需要一种快速的方法来计算形状为(45_000, 110, 12)的数据的jaccard_similarity那么我强烈建议您使用numbaparallel=True。我得到随机数据的运行时间形状(45_000, 110, 12)5 sec:(运行时获取colab)

import numpy as np
import numba as nb
import torch
@nb.jit(nopython=True, parallel=True)
def jaccard_Imahdi_Numba(batch_tensor):
tns_b = np.empty((batch_tensor.shape[0],batch_tensor.shape[1]-1))
for i in nb.prange(batch_tensor.shape[0]):
tns_tmp = batch_tensor[i][0]
for j , tns in enumerate(batch_tensor[i][1:]):
s1 = set(tns_tmp)
s2 = set(tns)
res = len(s1.intersection(s2)) / len(s1.union(s2))
tns_b[i][j] = res
return tns_b
large_tns = torch.tensor(np.random.randint(0,100, (45_000,110,12)))
%timeit jaccard_Imahdi_Numba(large_tns.numpy())
# 1 loop, best of 5: 5.37 s per loop

我在50_000 batch with shape (4,5) -> (50_000, 4, 5)的基准下面写。我们得到116 ms with numba但其他方法得到8 sec and 9 sec:(运行时获取colab)

import numpy as np
import numba as nb
import torch
def jaccard_asdf(batch_tensor):
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:
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)
a = torch.tensor(batch_tensor).numpy()  # just to demonstrate
ious = [jaccard(batch[:1, :], batch[1:, :]) for batch in a]
return np.array(ious).squeeze() 

def jaccard_Imahdi(batch_tensor):
def jaccard_similarity_numpy(list1, list2):
s1 = np.unique(list1)
s2 = np.unique(list2)
return (len(np.intersect1d(s1,s2))) / len(np.union1d(s1,s2))
tns_b = np.empty((batch_tensor.shape[0],batch_tensor.shape[1]-1))
for i in range(batch_tensor.shape[0]):
tns_tmp = batch_tensor[i][0]
for j , tns in enumerate(batch_tensor[i][1:]):
tns_b[i][j] = (jaccard_similarity_numpy(tns_tmp, tns))
return tns_b

@nb.jit(nopython=True, parallel=True)
def jaccard_Imahdi_Numba(batch_tensor):
tns_b = np.empty((batch_tensor.shape[0],batch_tensor.shape[1]-1))
for i in nb.prange(batch_tensor.shape[0]):
tns_tmp = batch_tensor[i][0]
for j , tns in enumerate(batch_tensor[i][1:]):
s1 = set(tns_tmp)
s2 = set(tns)
res = len(s1.intersection(s2)) / len(s1.union(s2))
tns_b[i][j] = res
return tns_b
small_tns = torch.tensor([
[[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]]
])
print(f'''output jaccard_asdf: n{
jaccard_asdf(small_tns)
}''')
print(f'''output jaccard_Imahdi: n{
jaccard_Imahdi(small_tns)
}''')
print(f'''output jaccard_Imahdi_Numba: n{
jaccard_Imahdi(small_tns)
}''')
large_tns = torch.tensor(np.random.randint(0,100, (50_000,4,5)))
%timeit jaccard_Imahdi(large_tns)
# 1 loop, best of 5: 8.32 s per loop
%timeit jaccard_asdf(large_tns)
# 1 loop, best of 5: 9.92 s per loop
%timeit jaccard_Imahdi_Numba(large_tns.numpy())
# 1 loop, best of 5: 116 ms per loop

输出:

output jaccard_asdf: 
[[0.28571429 0.4        0.16666667]
[0.14285714 0.16666667 0.14285714]]
output jaccard_Imahdi: 
[[0.28571429 0.4        0.16666667]
[0.14285714 0.16666667 0.14285714]]
output jaccard_Imahdi_Numba: 
[[0.28571429 0.4        0.16666667]
[0.14285714 0.16666667 0.14285714]]

最新更新