python实现数据集按比例划分

python实现数据集按比例划分

最近在做项目时,需要自己制作数据集,再把数据集划分为训练集和验证集。西瓜书上一般的划分比例是7:3或者8:2,这个代码闲来无事写了一下,比较容易。

import os
from sklearn import model_selection
def randomSplit(file):  #将存放数据集的路径传入
        sample = []
        trainSet = open(r'***.txt','w')  #打开你存放训练集名称的txt文件
        valSet = open(r'***.txt','w')    #打开存放验证机名称的txt文件
        for Image in os.listdir(file):
                sample.append(Image)
        #0.2表示验证集与训练集的比例为2:8,可以自行修改需要的比例。
        sample_train, sample_test = model_selection.train_test_split(sample, test_size=0.2) 
        for i in sample_train:
                trainSet.write(i)
                trainSet.write('\r\n') #加这句是为了输出每一个就换行
        for i in sample_test:
                valSet.write(i)
                valSet.write('\r\n')

代码运行结果:python实现数据集按比例划分_第1张图片

你可能感兴趣的:(计算机视觉,深度学习,python)