Python 实现数据集自动划分（训练集 / 验证集 / 测试集）

在深度学习模型训练中，我们通常需要将数据集划分为训练集（Train）、验证集（Val）和测试集（Test）。训练集用于模型参数学习，验证集用于超参数调优，测试集用于评估模型最终泛化能力。手动划分不仅效率低，还难以保证随机性，这里分享一段自动划分数据集的 Python 脚本。

代码功能说明

这段代码的核心功能是：将原始数据集中的图片和对应标签，按照 8:1:1 的比例随机划分为训练集、验证集和测试集，并分别存放于对应目录中，同时保证划分结果可复现。

代码如下

import os
import random
import shutil

# 分割训练集和验证集
# 设置随机种子以保证结果可复现
random.seed(42)

# 数据集根目录
data_dir = './datasets'
images_dir = os.path.join(data_dir, 'images')
labels_dir = os.path.join(data_dir, 'labels')

# 创建目录
os.makedirs(os.path.join(images_dir, 'train'), exist_ok=True)
os.makedirs(os.path.join(images_dir, 'val'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'train'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'val'), exist_ok=True)
os.makedirs(os.path.join(images_dir, 'test'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'test'), exist_ok=True)

# 获取所有的图像文件名
image_files = [f for f in os.listdir(os.path.join(data_dir, 'images')) if f.endswith('.jpg')]

# 计算训练集和验证集的大小
train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1
total_images = len(image_files)
train_index = int(total_images * train_ratio)
val_index = int(total_images * (train_ratio + val_ratio))

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

# 划分数据集
train_images = image_files[:train_index]
val_images = image_files[train_index:val_index]
test_images = image_files[val_index:]

for img_file in train_images:
    label_file = img_file.replace('.jpg', '.txt')
    shutil.copy(os.path.join(data_dir, 'images', img_file), os.path.join(images_dir, 'train'))
    shutil.copy(os.path.join(data_dir, 'labels', label_file), os.path.join(labels_dir, 'train'))

for img_file in val_images:
    label_file = img_file.replace('.jpg', '.txt')
    shutil.copy(os.path.join(data_dir, 'images', img_file), os.path.join(images_dir, 'val'))
    shutil.copy(os.path.join(data_dir, 'labels', label_file), os.path.join(labels_dir, 'val'))

for img_file in test_images:
    label_file = img_file.replace('.jpg', '.txt')
    shutil.copy(os.path.join(data_dir, 'images', img_file), os.path.join(images_dir, 'test'))
    shutil.copy(os.path.join(data_dir, 'labels', label_file), os.path.join(labels_dir, 'test'))

print(f"数据集划分完成！训练集：{len(train_images)}，验证集：{len(val_images)}，测试集：{len(test_images)}")

代码解析

1. 导入依赖库

首先导入必要的 Python 库，os 用于文件路径操作，random 用于随机打乱文件列表，shutil 用于文件复制：

import os
import random
import shutil

2. 保证结果可复现

设置随机种子，确保每次运行代码的划分结果一致（便于实验对比）：

random.seed(42)  # 固定种子，结果可复现

3. 目录设置与创建

定义数据集根目录及图片、标签存放路径，并自动创建划分后的子目录（train/val/test）：

# 数据集根目录
data_dir = './datasets'
images_dir = os.path.join(data_dir, 'images')  # 原始图片目录
labels_dir = os.path.join(data_dir, 'labels')  # 原始标签目录

# 创建划分后的子目录（若已存在则不报错）
os.makedirs(os.path.join(images_dir, 'train'), exist_ok=True)
os.makedirs(os.path.join(images_dir, 'val'), exist_ok=True)
os.makedirs(os.path.join(images_dir, 'test'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'train'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'val'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'test'), exist_ok=True)

4. 读取与划分文件

获取所有图片文件，按比例划分并随机打乱：

# 获取所有.jpg 格式的图片文件（可根据实际格式修改）
image_files = [f for f in os.listdir(images_dir) if f.endswith('.jpg')]

# 定义划分比例（可根据需求调整）
train_ratio = 0.8  # 训练集占比
val_ratio = 0.1    # 验证集占比
test_ratio = 0.1   # 测试集占比

# 计算各集合的文件数量
total_images = len(image_files)
train_index = int(total_images * train_ratio)      # 训练集结束索引
val_index = int(total_images * (train_ratio + val_ratio))  # 验证集结束索引

# 随机打乱文件列表（保证划分随机性）
random.shuffle(image_files)

# 划分数据集
train_images = image_files[:train_index]           # 训练集
val_images = image_files[train_index:val_index]    # 验证集
test_images = image_files[val_index:]              # 测试集

5. 复制文件到对应目录

将图片和对应的标签文件（假设标签与图片同名，后缀为.txt）复制到划分后的目录：

# 复制训练集文件
for img_file in train_images:
    label_file = img_file.replace('.jpg', '.txt')  # 标签文件名（与图片对应）
    shutil.copy(os.path.join(images_dir, img_file), os.path.join(images_dir, 'train'))
    shutil.copy(os.path.join(labels_dir, label_file), os.path.join(labels_dir, 'train'))

# 复制验证集文件（逻辑同上）
for img_file in val_images:
    label_file = img_file.replace('.jpg', '.txt')
    shutil.copy(os.path.join(images_dir, img_file), os.path.join(images_dir, 'val'))
    shutil.copy(os.path.join(labels_dir, label_file), os.path.join(labels_dir, 'val'))

# 复制测试集文件（逻辑同上）
for img_file in test_images:
    label_file = img_file.replace('.jpg', '.txt')
    shutil.copy(os.path.join(images_dir, img_file), os.path.join(images_dir, 'test'))
    shutil.copy(os.path.join(labels_dir, label_file), os.path.join(labels_dir, 'test'))

6. 输出划分结果

打印各集合的文件数量，确认划分完成：

print(f"数据集划分完成！训练集：{len(train_images)}，验证集：{len(val_images)}，测试集：{len(test_images)}")

使用说明

1. 格式适配：若图片格式为.png 等，需修改 endswith('.jpg') 为对应格式；标签格式不同时同理。
2. 比例调整：修改 train_ratio、val_ratio、test_ratio 可自定义划分比例。

目录结构要求：原始数据集需按如下结构存放（可修改代码中 data_dir 路径适配你的数据）：

datasets/
├─ images/  # 存放所有图片（.jpg 格式）
└─ labels/  # 存放所有标签（.txt 格式，与图片同名）