在观看 【好课推荐】深度学习与PyTorch入门实战 网络课程时,对于索引切片函数index_select()不是很理解,现分析如下:
import torch
# select by specific index
a = torch.rand(4, 3, 28, 28)
print(a.index_select(0, torch.tensor([0, 2])).shape) #1
#输出:torch.Size([2,3,28,28])
print(a.index_select(2, torch.arange(20)).shape) #2
#输出:torch.Size([4,3,20,28])
分析:
- a是一个4*3*28*28的四维张量,索引都是从0开始
-
index_select(dim,index) 参数说明: dim:表示从第几维挑选数据,类型为int值 index:表示从第一个参数维度中的哪个位置挑选数据,类型为torch.Tensor类的实例
- 对于1操作,从第0维(即数字“4”的的第一维)挑选数据,在该维度上选取索引值(编号)为第0和第2的列表(元组)的数据(共2组),其他地方不变,最后输出张量形状。#因此形状为:[4选2,3,28,28]
[2,3,28,28]
- 对于2操作,从第2维(即第一个数字“28”)挑选数据,在该维度上选取索引值(编号)为第0到第19的列表(元组)的数据(共20组),其他地方不变,最后输出张量形状。#因此形状为:[4,3,28选20,28]
[4,3,20,28]