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]最后生成的张量维度数减一。实际运用时也要多加注意。