将数据集划分为训练集与测试集

数据集划分

进行机器学习项目时我们经常会将数据集划分为训练集、测试集（和验证集）。为了保证数据的客观性，我们经常通过随机的方式打乱数据集，本文将提供一种随机分配数据集的方式。

环境

python 3

代码部分

进入数据集目录（默认所有图片存在同一文件夹内，如果储存在不同子文件夹内，请用os.walk()）
运行以下python代码。数据集会在此文件夹内分为三个文件夹：train, val, test

import os,shutil,random

val_rate=1
test_rate=2
train_rate=7
total_rate=val_rate+test_rate+train_rate

if not os.path.exists('train'):
    os.mkdir('train')
if not os.path.exists('test'):
    os.mkdir('test')
if not os.path.exists('val'):
    os.mkdir('val')
    
files = os.listdir()
files.remove('train')
files.remove('test')
files.remove('val')
lenf = len(files)
lenval = (lenf*val_rate)//total_rate
lentest = (lenf*test_rate)//total_rate
lentrain = lenf - lenval - lentest

indv = random.sample(range(lenf),lenval)
for i in indv:
    shutil.move(files[i],'val')
    
    #如果你需要同时移动标注(annotations)文件，假设你的标注文件为csv文件，
    #与图片存储在一起，要将其与图像一并移入val文件夹
    #shutil.move(files[i][:-4]+'.csv','val') 

files = os.listdir()
files.remove('train')
files.remove('test')
files.remove('val')
lenf = len(files)
indtest = random.sample(range(lenf),lentest)
for i in indtest:
    shutil.move(files[i],'test')#如果不想改变原位置数据，可以用shutil.copyfile()

files = os.listdir()
files.remove('train')
files.remove('test')
files.remove('val')
for i in files:
    shutil.move(i,'train')

代码为后期整理，如有bug感谢回复反馈，不定期更新