【pytorch】创建自己的数据集，并按类别的比例划分得到均匀的数据集

问题场景：当数据集有6个类别，每个类别的数据有150个，共900条数据，如图-1所示，每个文件夹中包含一类数据，共6个文件夹。想按照8：2的比例划分数据集，但又不想人为移动文件进行手动划分

图-1  数据

加载数据集使用datasets中的ImageFolder函数即可快速加载上图形式的数据。

代码：

from torchvision import datasets
Mydata = datasets.ImageFolder("./data",transform=torchvision.transforms.ToTensor())

问题：

当使用如下代码对数据进行划分时：

train_size = 120
test_size = 30
train_dataset, test_dataset  = random_split(Mydata,  lengths = [train_size, test_size], generator=torch.Generator().manual_seed(1))
# manual_seed(1)为随机数种子，保证每次的随机序列相同

会导致每一类的数据取得不均匀，也就是说，可能类别0取了30个，类别1取了26个，类别2取了38个这样的情况，我们希望每个类别都取30个。

#修改参数
classes = 6 
train_num = 120 
test_num = 30 
all_num = 150
#创建随机序列 
rand_list = [i for i in range(all_num)] 
random.seed(51) 
random.shuffle(rand_list) #打乱rand_list本身
train_list = rand_list[0:train_num] 
test_list = rand_list[train_num:all_num ] 
#采样的索引
train_index = [] 
test_index = [] 

for i in range(classes):     
    train_temp = list(map(lambda x: x + all_num * i,train_list))     
    train_index.extend(train_temp)     
    test_temp = list(map(lambda x: x + all_num * i, test_list))     
    test_index.extend(test_temp)     
train_dataset = Subset(Mydata,train_index) 
test_dataset = Subset(Mydata,test_index)