【深度学习小工具】手动划分训练集与验证集

前言

在研究深度学习时，经常会遇到一些未划分验证集的数据集，如果用训练集进行网络的验证，有可能会掩盖模型过拟合的现象，结果无说服力。

代码

import os
import shutil
import random

# train_ratio代表分割后训练集占原数据集的比率
def split_data(image_folder, label_folder, train_ratio=0.8):
    # 获取所有图片和标签文件的路径列表
    image_files = os.listdir(image_folder)
    label_files = os.listdir(label_folder)

    # 随机打乱文件列表
    random.shuffle(image_files)

    # 计算划分的训练集和验证集的边界索引
    split_index = int(len(image_files) * train_ratio)

    # 分割训练集和验证集
    train_images = image_files[:split_index]
    val_images = image_files[split_index:]

    # 创建训练集和验证集文件夹
    train_folder = 'train_data'
    val_folder = 'val_data'
    os.makedirs(train_folder, exist_ok=True)
    os.makedirs(val_folder, exist_ok=True)

    # 将训练集的图片和标签移动到训练集文件夹
    for image_file in train_images:
        shutil.move(os.path.join(image_folder, image_file), os.path.join(train_folder, image_file))
        corresponding_label = image_file.split('.')[0] + '.txt'
        shutil.move(os.path.join(label_folder, corresponding_label), os.path.join(train_folder, corresponding_label))

    # 将验证集的图片和标签移动到验证集文件夹
    for image_file in val_images:
        shutil.move(os.path.join(image_folder, image_file), os.path.join(val_folder, image_file))
        corresponding_label = image_file.split('.')[0] + '.txt'
        shutil.move(os.path.join(label_folder, corresponding_label), os.path.join(val_folder, corresponding_label))

# 用法示例
image_folder_path = 'path/to/your/image_folder'
label_folder_path = 'path/to/your/label_folder'
split_data(image_folder_path, label_folder_path, train_ratio=0.8)

使用时只需按需调整train_ratio的数字即可实现训练集与验证集的划分。
划分后的训练集（验证集）的图片和标签在同一个文件夹，可以在终端用命令快速移动到所需的文件夹，以ubuntu为例：

mv *.txt /path/