pytorch 如何把圖像數(shù)據(jù)集進(jìn)行劃分成train,test和val

更新時(shí)間：2021年05月31日 10:42:36 作者：l8947943

這篇文章主要介紹了pytorch 把圖像數(shù)據(jù)集進(jìn)行劃分成train,test和val的操作，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。如有錯(cuò)誤或未考慮完全的地方，望不吝賜教

1、手上目前擁有數(shù)據(jù)集是一大坨，沒(méi)有train，test，val的劃分

如圖所示

在這里插入圖片描述

2、目錄結(jié)構(gòu)：

|---data
     |---dslr
         |---images
         		|---back_pack
         			|---a.jpg
         			|---b.jpg
         			...

3、轉(zhuǎn)換后的格式如圖

在這里插入圖片描述

目錄結(jié)構(gòu)為：

|---datanews
     |---dslr
         |---images
         		|---test
         		|---train
         		|---valid
	         		|---back_pack
	         			|---a.jpg
	         			|---b.jpg
	         			...

4、代碼如下：

4.1 先創(chuàng)建同樣結(jié)構(gòu)的層級(jí)結(jié)構(gòu)

4.2 然后講原始數(shù)據(jù)按照比例劃分

4.3 移入到對(duì)應(yīng)的文件目錄里面

import os, random, shutil

def make_dir(source, target):
    '''
    創(chuàng)建和源文件相似的文件路徑函數(shù)
    :param source: 源文件位置
    :param target: 目標(biāo)文件位置
    '''
    dir_names = os.listdir(source)
    for names in dir_names:
        for i in ['train', 'valid', 'test']:
            path = target + '/' + i + '/' + names
            if not os.path.exists(path):
                os.makedirs(path)

def divideTrainValiTest(source, target):
    '''
        創(chuàng)建和源文件相似的文件路徑
        :param source: 源文件位置
        :param target: 目標(biāo)文件位置
    '''
    # 得到源文件下的種類
    pic_name = os.listdir(source)
    
    # 對(duì)于每一類里的數(shù)據(jù)進(jìn)行操作
    for classes in pic_name:
        # 得到這一種類的圖片的名字
        pic_classes_name = os.listdir(os.path.join(source, classes))
        random.shuffle(pic_classes_name)
        
        # 按照8：1：1比例劃分
        train_list = pic_classes_name[0:int(0.8 * len(pic_classes_name))]
        valid_list = pic_classes_name[int(0.8 * len(pic_classes_name)):int(0.9 * len(pic_classes_name))]
        test_list = pic_classes_name[int(0.9 * len(pic_classes_name)):]
        
        # 對(duì)于每個(gè)圖片，移入到對(duì)應(yīng)的文件夾里面
        for train_pic in train_list:
            shutil.copyfile(source + '/' + classes + '/' + train_pic, target + '/train/' + classes + '/' + train_pic)
        for validation_pic in valid_list:
            shutil.copyfile(source + '/' + classes + '/' + validation_pic,
                            target + '/valid/' + classes + '/' + validation_pic)
        for test_pic in test_list:
            shutil.copyfile(source + '/' + classes + '/' + test_pic, target + '/test/' + classes + '/' + test_pic)

if __name__ == '__main__':
    filepath = r'../data/dslr/images'
    dist = r'../datanews/dslr/images'
    make_dir(filepath, dist)
    divideTrainValiTest(filepath, dist)

補(bǔ)充：pytorch中數(shù)據(jù)集的劃分方法及eError: take(): argument 'index' (position 1) must be Tensor, not numpy.ndarray錯(cuò)誤原因

在使用pytorch框架時(shí)，難免需要對(duì)數(shù)據(jù)集進(jìn)行訓(xùn)練集和驗(yàn)證集的劃分，一般使用sklearn.model_selection中的train_test_split方法

該方法使用如下：

from sklearn.model_selection import train_test_split
import numpy as np
import torch
import torch.autograd import Variable
from torch.utils.data import DataLoader
 
traindata = np.load(train_path)   # image_num * W * H
trainlabel = np.load(train_label_path)
train_data = traindata[:, np.newaxis, ...]
train_label_data = trainlabel[:, np.newaxis, ...]
 
x_tra, x_val, y_tra, y_val = train_test_split(train_data, train_label_data, test_size=0.1, random_state=0)  # 訓(xùn)練集和驗(yàn)證集使用9:1
 
x_tra = Variable(torch.from_numpy(x_tra))
x_tra = x_tra.float()
y_tra = Variable(torch.from_numpy(y_tra))
y_tra = y_tra.float()
 
x_val = Variable(torch.from_numpy(x_val))
x_val = x_val.float()
y_val = Variable(torch.from_numpy(y_val))
y_val = y_val.float()
 
# 訓(xùn)練集的DataLoader
traindataset = torch.utils.data.TensorDataset(x_tra, y_tra)
trainloader = DataLoader(dataset=traindataset, num_workers=opt.threads, batch_size=8, shuffle=True)  
 
# 驗(yàn)證集的DataLoader
validataset = torch.utils.data.TensorDataset(x_val, y_val)
valiloader = DataLoader(dataset=validataset, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)

注意：如果按照如下方式使用，就會(huì)報(bào)eError: take(): argument 'index' (position 1) must be Tensor, not numpy.ndarray錯(cuò)誤

from sklearn.model_selection import train_test_split
import numpy as np
import torch
import torch.autograd import Variable
from torch.utils.data import DataLoader
 
traindata = np.load(train_path)   # image_num * W * H
trainlabel = np.load(train_label_path)
 
train_data = traindata[:, np.newaxis, ...]
train_label_data = trainlabel[:, np.newaxis, ...]
 
x_train = Variable(torch.from_numpy(train_data))
x_train = x_train.float()
y_train = Variable(torch.from_numpy(train_label_data))
y_train = y_train.float()
# 將原始的訓(xùn)練數(shù)據(jù)集分為訓(xùn)練集和驗(yàn)證集，后面就可以使用早停機(jī)制
x_tra, x_val, y_tra, y_val = train_test_split(x_train, y_train, test_size=0.1)  # 訓(xùn)練集和驗(yàn)證集使用9:1