Pytorch 对三维向量进行随机采样

我输入一个大小为torch.Size([2, 6, 5])的三维向量，现在我要从中随机选取N个点，且为不重复采样。为提供更加清楚的解释，我们运行结果如下截图可做对比：

1. 无序采样

随机采样输出是没有顺序的，可能在位置2的点采样后位置在3！
a) 从第二维度进行随机采样

b) 从第三维度进行随机采样

2. 有序采样

采样输出跟原本数据的先后顺序一致

3. 代码

import torch
import random

# 输入大小为torch.Size([2, 6, 5])的三维向量
a = torch.randn(2,6,5)
N=4  # 采样数目

# 第二维随机采样点下标
index=torch.LongTensor(random.sample(range(a.size(1)),N)) 
# 第三维随机采样点下标
index=torch.LongTensor(random.sample(range(a.size(2)),N)) 
# 第M维随机采样后输出点的先后顺序不变
index=torch.LongTensor(random.sample(range(a.size(M-1)),N)).sort()[0]

# 第二维随机采样后的输出，这里输出torch.Size([2, 4, 5])
b=torch.index_select(a,1,index)  
# 第三维随机采样后的输出，这里输出torch.Size([2, 6, 4])
b=torch.index_select(a,2,index)

借鉴博客：Pytorch 对张量进行随机采样

看到这么努力的份上，点个赞再走吧