PyTorch Tensor对Tensor索引（Index wrt Tensor）


简介

模型越编越复杂，有时候就会涉及到一些PyTorch里面没有的Tensor操作，这时候就得自己写。

最近需要实现Tensor对Tensor的索引，即对于一个Tensor中的每个元素，返回它在另一个Tensor中的位置，例子如下：

1
2
3
4
5
> inp = torch.LongTensor([5, 3, 6, 9])
> index = torch.LongTensor([2, 3, 4, 6, 8, 5, 7, 9])
> out = tensor_index(inp, index)
> print(out)
tensor([5, 1, 3, 7])

找了一圈，似乎没有实现此功能的库函数，于是自己实现一个。

基础实现

实现思路： 将输入向量增加一个纬度后，与索引内的每个元素判断相等，然后收集相等的位置的信息，作为结果。

1
2
3
4
5
6
7
8
9
10
> x = torch.LongTensor([5, 3, 6, 9])
> y = torch.LongTensor([2, 3, 4, 6, 8, 5, 7, 9])
> r = x.unsqueeze(-1) == y
> print(r)
tensor([[False, False, False, False, False,  True, False, False],
        [False,  True, False, False, False, False, False, False],
        [False, False, False,  True, False, False, False, False],
        [False, False, False, False, False, False, False,  True]])
> r.long().argmax(dim=-1)
tensor([5, 1, 3, 7])

完善

作为一个成熟的功能函数，还得要考虑各种情况并做出处理，比如对index中没有的元素索引怎么办，index的纬度不是1怎么办，等等。以下是一个完善后的函数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def tensor_index(inp, index):
    # 要求index的维度必须为1
    assert index.dim() == 1, "The dimension of index must be 1."
    # 查找index中的重复元素
    index_set = set(index.tolist())
    assert index.shape[0] == len(index_set), "There are repeated element(s) in index Tensor."

    r = inp.unsqueeze(-1) == index
    r = r.long()

    # 处理索引不到的情况
    sum_check = r.sum(dim=-1)
    sum_check_0 = sum_check == 0
    assert not sum_check_0.any(), "Some element(s) in inp Tensor cannot be indexed in index Tensor."
   
    r = r.argmax(dim=-1)

    return r

参考

  • https://discuss.pytorch.org/t/find-indices-of-one-tensor-in-another/84889/2