欢迎光临散文网 会员登陆 & 注册

Pytorch将数据集划分为训练集、验证集和测试集

2020-08-26 16:46 作者:肆十二-  | 我要投稿

Pytorch将数据集划分为训练集、验证集和测试集

我们可以借助Pytorch从文件夹中读取数据集,十分方便,但是Pytorch中没有提供数据集划分的操作,需要手动将原始的数据集划分为训练集、验证集和测试集,废话不多说,这里我写了一个工具类,帮助大家将数据集自动划分为训练集、验证集和测试集,还可以指定比例,代码如下。

# 工具类
import os
import random
import shutil
from shutil import copy2


def data_set_split(src_data_folder, target_data_folder, train_scale=0.8, val_scale=0.1, test_scale=0.1):
   '''
   读取源数据文件夹,生成划分好的文件夹,分为trian、val、test三个文件夹进行
   :param src_data_folder: 源文件夹 E:/biye/gogogo/note_book/torch_note/data/utils_test/data_split/src_data
   :param target_data_folder: 目标文件夹 E:/biye/gogogo/note_book/torch_note/data/utils_test/data_split/target_data
   :param train_scale: 训练集比例
   :param val_scale: 验证集比例
   :param test_scale: 测试集比例
   :return:
   '''
   print("开始数据集划分")
   class_names = os.listdir(src_data_folder)
   # 在目标目录下创建文件夹
   split_names = ['train', 'val', 'test']
   for split_name in split_names:
       split_path = os.path.join(target_data_folder, split_name)
       if os.path.isdir(split_path):
           pass
       else:
           os.mkdir(split_path)
       # 然后在split_path的目录下创建类别文件夹
       for class_name in class_names:
           class_split_path = os.path.join(split_path, class_name)
           if os.path.isdir(class_split_path):
               pass
           else:
               os.mkdir(class_split_path)

   # 按照比例划分数据集,并进行数据图片的复制
   # 首先进行分类遍历
   for class_name in class_names:
       current_class_data_path = os.path.join(src_data_folder, class_name)
       current_all_data = os.listdir(current_class_data_path)
       current_data_length = len(current_all_data)
       current_data_index_list = list(range(current_data_length))
       random.shuffle(current_data_index_list)

       train_folder = os.path.join(os.path.join(target_data_folder, 'train'), class_name)
       val_folder = os.path.join(os.path.join(target_data_folder, 'val'), class_name)
       test_folder = os.path.join(os.path.join(target_data_folder, 'test'), class_name)
       train_stop_flag = current_data_length * train_scale
       val_stop_flag = current_data_length * (train_scale + val_scale)
       current_idx = 0
       train_num = 0
       val_num = 0
       test_num = 0
       for i in current_data_index_list:
           src_img_path = os.path.join(current_class_data_path, current_all_data[i])
           if current_idx <= train_stop_flag:
               copy2(src_img_path, train_folder)
               # print("{}复制到了{}".format(src_img_path, train_folder))
               train_num = train_num + 1
           elif (current_idx > train_stop_flag) and (current_idx <= val_stop_flag):
               copy2(src_img_path, val_folder)
               # print("{}复制到了{}".format(src_img_path, val_folder))
               val_num = val_num + 1
           else:
               copy2(src_img_path, test_folder)
               # print("{}复制到了{}".format(src_img_path, test_folder))
               test_num = test_num + 1

           current_idx = current_idx + 1

       print("*********************************{}*************************************".format(class_name))
       print(
           "{}类按照{}:{}:{}的比例划分完成,一共{}张图片".format(class_name, train_scale, val_scale, test_scale, current_data_length))
       print("训练集{}:{}张".format(train_folder, train_num))
       print("验证集{}:{}张".format(val_folder, val_num))
       print("测试集{}:{}张".format(test_folder, test_num))


if __name__ == '__main__':
   src_data_folder = "E:/biye/gogogo/note_book/torch_note/data/utils_test/data_split/src_data"
   target_data_folder = "E:/biye/gogogo/note_book/torch_note/data/utils_test/data_split/target_data"
   data_set_split(src_data_folder, target_data_folder)

** 注意 **

划分前你得文件夹结构应该是这样的

image-20200826160553697

划分结果

data_split

tensorflow2.3 加载数据集的方式

from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential


def load_data_from_folder(batch_size, target_img_height, target_img_width, data_dir="F:/datas/massmass/fer2013+/ccc/"):
   train_datagen = ImageDataGenerator(
       rescale=1. / 255,  # 重放缩因子,数值乘以1.0/255(归一化)
       shear_range=0.2,  # 剪切强度(逆时针方向的剪切变换角度)
       zoom_range=0.2,  # 随机缩放的幅度
       # 进行随机水平翻转
       horizontal_flip=True)
   val_datagen = ImageDataGenerator(
       rescale=1. / 255)

   train_generator = train_datagen.flow_from_directory(
       data_dir + '/train',  # dictory参数,该路径下的所有子文件夹的图片都会被生成器使用,无限产生batch数据
       target_size=(target_img_height, target_img_width),  # 图片将被resize成该尺寸
       color_mode='grayscale',  # 颜色模式,graycsale或rgb(默认rgb)
       batch_size=batch_size,  # batch数据的大小,默认为32
       class_mode='sparse')  # 返回的标签形式,默认为‘category’,返回2D的独热码标签
   val_generator = val_datagen.flow_from_directory(
       data_dir + '/val',  # 同上
       target_size=(target_img_height, target_img_width),
       color_mode='grayscale',
       batch_size=batch_size,
       class_mode='sparse')
   num_class = train_generator.num_classes
   return train_generator, val_generator, num_class

tensorflow2.0 加载数据集的方式

from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential

def load_data_from_folder(batch_size, target_img_height, target_img_width, data_dir="data/"):
   emotion_classification_train_datagen = ImageDataGenerator(
       rescale=1. / 255,  # 重放缩因子,数值乘以1.0/255(归一化)
       shear_range=0.2,  # 剪切强度(逆时针方向的剪切变换角度)
       zoom_range=0.2,  # 随机缩放的幅度
       # 进行随机水平翻转
       horizontal_flip=True)
   emotion_classification_val_datagen = ImageDataGenerator(
       rescale=1. / 255)

   emotion_classification_train_generator = emotion_classification_train_datagen.flow_from_directory(
       data_dir + '/train',  # dictory参数,该路径下的所有子文件夹的图片都会被生成器使用,无限产生batch数据
       target_size=(target_img_height, target_img_width),  # 图片将被resize成该尺寸
       color_mode='grayscale',  # 颜色模式,graycsale或rgb(默认rgb)
       batch_size=batch_size,  # batch数据的大小,默认为32
       class_mode='sparse')  # 返回的标签形式,默认为‘category’,返回2D的独热码标签
   emotion_classification_val_generator = emotion_classification_val_datagen.flow_from_directory(
       data_dir + '/val',  # 同上
       target_size=(target_img_height, target_img_width),
       color_mode='grayscale',
       batch_size=batch_size,
       class_mode='sparse')
   num_class = emotion_classification_train_generator.num_classes
   return emotion_classification_train_generator, emotion_classification_val_generator, num_class


pytorch加载数据集的方式

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy

# Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {
   'train': transforms.Compose([
       transforms.RandomResizedCrop(224),
       transforms.RandomHorizontalFlip(),
       transforms.ToTensor(),
       transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
   ]),
   'val': transforms.Compose([
       transforms.Resize(256),
       transforms.CenterCrop(224),
       transforms.ToTensor(),
       transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
   ]),
}

data_dir = 'data/hymenoptera_data'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                         data_transforms[x])
                 for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                            shuffle=True, num_workers=4)
             for x in ['train', 'val']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


最后附上github地址

https://github.com/cmFighting/mnist_demo_torch1.6



Pytorch将数据集划分为训练集、验证集和测试集的评论 (共 条)

分享到微博请遵守国家法律