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

Logo

华为开发者空间,是为全球开发者打造的专属开发空间,汇聚了华为优质开发资源及工具,致力于让每一位开发者拥有一台云主机,基于华为根生态开发、创新。

更多推荐