pytorch 中二维张量的高级索引

Question

我有一个二维张量 X. 和两个索引列表，第一个索引和第二个调用 a 和 b。 我想做

X[a[i],b[i]] = 0 for i in range(len(a))

我怎么能做到这一点。 如果我直接做X[a,b]错误是 IndexError: The advanced indexing objects could not be broadcast

Answer 1

检查包含索引的lists ，某些值可能超出范围。 那时你会得到如下所示的IndexError ：

在 [43] 中：X[4,4]

IndexError Traceback（最近一次调用最后一次） in () ----> 1 X[4,4]

索引错误：索引 4 超出维度 0（大小为 3）的范围

如果您的指数在正确的范围内，它应该可以正常工作。

下面是一个例子：

In [35]: X = torch.Tensor([[3, 4, 5, 6], [1, 2, 3, 4], [6, 3, 2, 1]])

In [36]: X
Out[36]: 

 3  4  5  6
 1  2  3  4
 6  3  2  1
[torch.FloatTensor of size 3x4]

In [37]: a = [0, 2]

In [38]: b = [1, 2]

In [39]: X[a, b]
Out[39]: 

 4
 2
[torch.FloatTensor of size 2]

In [40]: X[a, b] = 0

In [41]: X
Out[41]: 

 3  0  5  6
 1  2  3  4
 6  3  0  1
[torch.FloatTensor of size 3x4]