pytorch中index_select()的详细用法与注意事项

index_select()

index_select()函数原型为：

 def index_select(self, dim: _int, index: Tensor) -> Tensor: ...

第一参数大家都知道是表示对象。

第二个参数dim表示维度，具体取哪一个维度就要看你对象原本的维度数了，比如一个torch.Size([1, 3, 3])的张量，你想取某一列，那dim就该取2或-1，一个torch.Size([ 3,3])的张量，想取某一行的话，那dim就该取1或-1。

第三个参数index表示你想取的索引的序号，比如torch.tensor([0, 1])就是取第一第二列。

但是要注意，最后生成的张量,其维度是保持不变的！下面举个例子：

a = torch.Tensor([[[1,2,3],
                   [4,5,6],
                   [7,8,9]]])

b = a.index_select(2,torch.LongTensor([0]))
c = a[:,:,0]
print("a.shape:{}".format(a.shape))
print("b:{} \n b.shape:{}".format(b,b.shape))
print("c:{} \n c.shape:{}".format(c,c.shape))

结果为：

a.shape:torch.Size([1, 3, 3])
b:tensor([[[1.],
         [4.],
         [7.]]]) 
 b.shape:torch.Size([1, 3, 1])
c:tensor([[1., 4., 7.]]) 
 c.shape:torch.Size([1, 3])

可以看到通过index_select（）和a[:,:,0]的方式都取出了第一列的数据，但是index_select最后生成的张量是维度不变的，而a[:,:,0]最后生成的张量维度数减一。实际运用时也要多加注意。