在pytorch中索引支持类似numpy方式,还有一种是narrow方式。
区别直接看代码:
import torch
data=torch.tensor([[1,2],[3,4],[5,6]])
#tensor([[1, 2],
[3, 4],
[5, 6]])
#使用numpy方式从data中取出第二列数据
In [40]: data[:,1]
Out[40]: tensor([2, 4, 6])
In [41]: data[:,1].size()
Out[41]: torch.Size([3])
#使用narrow方式取出第二列数据
In [46]: data.narrow(1,1,1)
Out[46]:
tensor([[2],
[4],
[6]])
In [47]: data.narrow(1,1,1).size()
Out[47]: torch.Size([3, 1])
可以看出narrow返回数据还是原数据的格式
第一个参数表示我是对列操作还是对行操作
0表示我要对行操作
1表示我要对列操作
比如我要取第二列的数据如下,第二个参数表示开始位置,第三个参数是取的个数(0,1,1)表示对行进行操作,从第2行开始(索引从0开始),取1个。
In [48]: data.narrow(0,1,1)
Out[48]: tensor([[3, 4]])