Python實(shí)現(xiàn)數(shù)據(jù)集自動(dòng)劃分的示例代碼(訓(xùn)練集/驗(yàn)證集/測試集)

更新時(shí)間：2026年04月02日 09:41:16 作者：CV小涵

本文提供了一段Python腳本,用于自動(dòng)將數(shù)據(jù)集按照指定比例劃分為訓(xùn)練集、驗(yàn)證集和測試集,并保證結(jié)果可復(fù)現(xiàn),腳本涵蓋了目錄設(shè)置、文件讀取與劃分、文件復(fù)制及結(jié)果輸出等步驟,適用于深度學(xué)習(xí)模型的訓(xùn)練與評(píng)估,需要的朋友可以參考下

在深度學(xué)習(xí)模型訓(xùn)練中，我們通常需要將數(shù)據(jù)集劃分為訓(xùn)練集（Train）、驗(yàn)證集（Val）和測試集（Test）。訓(xùn)練集用于模型參數(shù)學(xué)習(xí)，驗(yàn)證集用于超參數(shù)調(diào)優(yōu)，測試集用于評(píng)估模型最終泛化能力。手動(dòng)劃分不僅效率低，還難以保證隨機(jī)性，這里分享一段自動(dòng)劃分?jǐn)?shù)據(jù)集的 Python 腳本。

代碼功能說明

這段代碼的核心功能是：將原始數(shù)據(jù)集中的圖片和對(duì)應(yīng)標(biāo)簽，按照 8:1:1 的比例隨機(jī)劃分為訓(xùn)練集、驗(yàn)證集和測試集，并分別存放于對(duì)應(yīng)目錄中，同時(shí)保證劃分結(jié)果可復(fù)現(xiàn)。

代碼如下

import os
import random
import shutil
# 分割訓(xùn)練集和驗(yàn)證集
# 設(shè)置隨機(jī)種子以保證結(jié)果可復(fù)現(xiàn)
random.seed(42)

# 數(shù)據(jù)集根目錄
data_dir = './datasets'

images_dir = os.path.join(data_dir, 'images')
labels_dir = os.path.join(data_dir, 'labels')

# 創(chuàng)建目錄
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')]

# 計(jì)算訓(xùn)練集和驗(yàn)證集的大小
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)

# 劃分?jǐn)?shù)據(jù)集
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"數(shù)據(jù)集劃分完成！訓(xùn)練集: {len(train_images)}，驗(yàn)證集: {len(val_images)}，測試集: {len(test_images)}")
import sys
sys.exit(0)

代碼解析

1. 導(dǎo)入依賴庫

首先導(dǎo)入必要的 Python 庫，os用于文件路徑操作，random用于隨機(jī)打亂文件列表，shutil用于文件復(fù)制：

import os
import random
import shutil

2. 保證結(jié)果可復(fù)現(xiàn)

設(shè)置隨機(jī)種子，確保每次運(yùn)行代碼的劃分結(jié)果一致（便于實(shí)驗(yàn)對(duì)比）：

random.seed(42)  # 固定種子，結(jié)果可復(fù)現(xiàn)

3. 目錄設(shè)置與創(chuàng)建

定義數(shù)據(jù)集根目錄及圖片、標(biāo)簽存放路徑，并自動(dòng)創(chuàng)建劃分后的子目錄（train/val/test）：

# 數(shù)據(jù)集根目錄
data_dir = './datasets'
images_dir = os.path.join(data_dir, 'images')  # 原始圖片目錄
labels_dir = os.path.join(data_dir, 'labels')  # 原始標(biāo)簽?zāi)夸?

# 創(chuàng)建劃分后的子目錄（若已存在則不報(bào)錯(cuò)）
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. 讀取與劃分文件

獲取所有圖片文件，按比例劃分并隨機(jī)打亂：

# 獲取所有.jpg格式的圖片文件（可根據(jù)實(shí)際格式修改）
image_files = [f for f in os.listdir(images_dir) if f.endswith('.jpg')]

# 定義劃分比例（可根據(jù)需求調(diào)整）
train_ratio = 0.8  # 訓(xùn)練集占比
val_ratio = 0.1    # 驗(yàn)證集占比
test_ratio = 0.1   # 測試集占比

# 計(jì)算各集合的文件數(shù)量
total_images = len(image_files)
train_index = int(total_images * train_ratio)  # 訓(xùn)練集結(jié)束索引
val_index = int(total_images * (train_ratio + val_ratio))  # 驗(yàn)證集結(jié)束索引

# 隨機(jī)打亂文件列表（保證劃分隨機(jī)性）
random.shuffle(image_files)

# 劃分?jǐn)?shù)據(jù)集
train_images = image_files[:train_index]       # 訓(xùn)練集
val_images = image_files[train_index:val_index]  # 驗(yàn)證集
test_images = image_files[val_index:]          # 測試集

5. 復(fù)制文件到對(duì)應(yīng)目錄

將圖片和對(duì)應(yīng)的標(biāo)簽文件（假設(shè)標(biāo)簽與圖片同名，后綴為.txt）復(fù)制到劃分后的目錄：

# 復(fù)制訓(xùn)練集文件
for img_file in train_images:
    label_file = img_file.replace('.jpg', '.txt')  # 標(biāo)簽文件名（與圖片對(duì)應(yīng)）
    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'))

# 復(fù)制驗(yàn)證集文件（邏輯同上）
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'))

# 復(fù)制測試集文件（邏輯同上）
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. 輸出劃分結(jié)果

打印各集合的文件數(shù)量，確認(rèn)劃分完成：

print(f"數(shù)據(jù)集劃分完成！訓(xùn)練集: {len(train_images)}，驗(yàn)證集: {len(val_images)}，測試集: {len(test_ima

使用說明

目錄結(jié)構(gòu)要求：原始數(shù)據(jù)集需按如下結(jié)構(gòu)存放（可修改代碼中data_dir路徑適配你的數(shù)據(jù)）：

datasets/
├─ images/  # 存放所有圖片（.jpg格式）
└─ labels/  # 存放所有標(biāo)簽（.txt格式，與圖片同名）

格式適配：若圖片格式為.png 等，需修改endswith('.jpg')為對(duì)應(yīng)格式；標(biāo)簽格式不同時(shí)同理。
比例調(diào)整：修改train_ratio、val_ratio、test_ratio可自定義劃分比例。

以上就是Python實(shí)現(xiàn)數(shù)據(jù)集自動(dòng)劃分的示例代碼(訓(xùn)練集/驗(yàn)證集/測試集)的詳細(xì)內(nèi)容，更多關(guān)于Python數(shù)據(jù)集自動(dòng)劃分的資料請(qǐng)關(guān)注腳本之家其它相關(guān)文章！

您可能感興趣的文章: