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]