sklearn 划分数据集

Cincinnati_De

于 2018-09-13 18:44:17 发布

阅读量1.5k

点赞数

CC 4.0 BY-SA版权

分类专栏：个人日记

本文链接：https://blog.youkuaiyun.com/Cincinnati_De/article/details/80622948

个人日记专栏收录该内容

142 篇文章

订阅专栏

本文介绍如何使用Sklearn库中的train_test_split函数来划分数据集为训练集和测试集，通过设置测试集比例、随机状态和分层抽样参数，确保数据划分的合理性和有效性。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

#-*- coding: UTF-8 -*-

from sklearn.model_selection import train_test_split



def split(dataset, labelset, test_size, train_savefile, test_savefile):

    # split into training set and test set
    x_train, x_test, y_train, y_test = train_test_split(dataset, labelset, test_size=test_size, random_state=42, stratify=labelset )
  
    savetxt(train_savefile, x_train)
    savetxt(test_savefile, x_test)

    return x_train, x_test


def savetxt(path, np_array):
    with open(file=path, mode='w', encoding='utf-8') as fw:
        fw.writelines(np_array)

def reader_data(datafile):
    data_list = []
    with open(datafile, mode='r', encoding='utf-8') as f:
        for line in f:
            data_list.append(line)

    return data_list

if __name__ == '__main__':

    datafile = 'data/output/tra-set0603_0.9'
    dataset = reader_data(datafile)
    label_file = 'data/output/tra-set0603_0.9_label'
    labelset = reader_data(label_file)

    test_size = 0.2
    train_savefile = 'data/output/raw_0.9/raw_train.txt'
    test_savefile = 'data/output/raw_0.9/raw_test.txt'
    split(dataset, labelset, test_size, train_savefile, test_savefile)