PyTorch中index_select()函数的理解

在观看 【好课推荐】深度学习与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])

分析：

1. a是一个4*3*28*28的四维张量，索引都是从0开始

2. 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]

 课程链接：【好课推荐】深度学习与PyTorch入门实战_哔哩哔哩_bilibili