例子
a = torch.linspace(1, 12, steps=12).view(3, 4) print(a) b = torch.index_select(a, 0, torch.tensor([0, 2])) print(b) print(a.index_select(0, torch.tensor([0, 2]))) c = torch.index_select(a, 1, torch.tensor([1, 3])) print(c)
先定義了一個tensor,這里用到了linspace和view方法。
第一個參數是索引的對象,第二個參數0表示按行索引,1表示按列進行索引,第三個參數是一個tensor,就是索引的序號,比如b里面tensor[0, 2]表示第0行和第2行,c里面tensor[1, 3]表示第1列和第3列。
輸出結果如下:
tensor([[ 1., 2., 3., 4.], [ 5., 6., 7., 8.], [ 9., 10., 11., 12.]]) tensor([[ 1., 2., 3., 4.], [ 9., 10., 11., 12.]]) tensor([[ 1., 2., 3., 4.], [ 9., 10., 11., 12.]]) tensor([[ 2., 4.], [ 6., 8.], [10., 12.]])