pytorch中index_select()的用法
程序员文章站
2022-04-18 20:06:05
...
pytorch中index_select()的用法
刚开始学习pytorch,遇到了index_select(),一开始不太明白几个参数的意思,后来查了一下资料,算是明白了一点。
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.]])