时间戳

首页 > 解决方案 > 如何使用 sklearn 将数据集拆分为训练和测试数据

image-processing - 如何使用 sklearn 将数据集拆分为训练和测试数据

问题描述

我正在使用 at&t faces 数据集，一个主目录包含 40 个子目录，每个子目录包含特定人的不同图像。我创建了一个包含子目录名称的列表。我想用这些数据来训练一个神经网络，所以我想把数据分成 80% 的训练和 20% 的测试。这是我到目前为止所做的：

import os
import cv2
path = r"C:\Users\Desktop\att_faces"
directory = []
directory = [x[1] for x in os.walk(path)]
non_empty_dirs = [x for x in directory if x]


directory = [item for subitem in non_empty_dirs for item in subitem]

这一步之后我应该如何进行？

标签： image-processingmachine-learningscikit-learndeep-learningcomputer-vision

解决方案

您希望将数据拆分为训练集和测试集。为此，您可以

  1. 手动或使用脚本单独训练和测试文件夹，并在数据生成器的帮助下加载它们进行训练。

  2. 加载整个数据并将它们拆分以在内存中进行训练和测试。

让我们讨论第二个选项。

主目录包含40个子目录

假设您的主目录是Train//并且有40子文件夹即1-40. 另外，我假设类标签是文件夹名称。

# imports

import cv2
import numpy as np
import os

from sklearn.model_selection import train_test_split

# seed for reproducibility
SEED = 44000


# lists to store data
data = []
label = []

# folder where data is placed
BASE_FOLDER = 'Train//'

folders = os.listdir(BASE_FOLDER)

# loading data to lists
for folder in folders:
    for file in os.listdir(BASE_FOLDER + folder + '//'):
        img = cv2.imread(BASE_FOLDER + folder + '//' + file)
        # do any pre-processing if needed like resize, sharpen etc.
        data = data.append(img)
        label = label.append(folder)

# now split the data in to train and test with the help of train_test_split

train_data, test_data, train_label, test_label = train_test_split(data, label, test_size=0.2, random_state=SEED)

推荐阅读