Python 训练集、测试集以及验证集切分方法：sklearn及手动切分

目录

方法一

方法二

需求目的：针对模型训练输入，按照6:2:2的比例进行训练集、测试集和验证集的划分。当前数据量约10万条。如果针对的是记录条数达上百万的数据集，可按照98:1:1的比例进行切分。

方法一：切分训练集和测试集，采用机器学习包sklearn中的train_test_split()函数
方法二：切分训练集、测试集以及验证集，针对dataframe手动切分

方法一

采用Sklearn包中的sklearn.model_selection.train_test_split()函数，该函数功能是将原始数据按照比例切分为训练集和测试集。

1. 函数形式：
2. sklearn.model_selection.train_test_split(*arrays, test_size=None,
3. train_size=None, random_state=None, shuffle=True, stratify=None)
5. 参数解读：
6. *arrays：等长的列表、数组或者dataframe等
7. test_size: 0和1之间，默认0.25
8. train_size: 0和1之间，默认1
9. random_state: 传递一个int值，以便在多个函数调用之间产生可复制的输出
10. shuffle: 拆分前是否进行洗牌
11. strafity: 是否对数据进行分层
13. 返回结果：
14. 输入序列的train test分割序列

例子

1. >>> import numpy as np
2. >>> from sklearn.model_selection import train_test_split
3. >>> X, y = np.arange(10).reshape((5, 2)), range(5)
4. >>> X
5. array([[0, 1],
6. [2, 3],
7. [4, 5],
8. [6, 7],
9. [8, 9]])
10. >>> list(y)
11. [0, 1, 2, 3, 4]
13. >>> X_train, X_test, y_train, y_test = train_test_split(
14. ... X, y, test_size=0.33, random_state=42)
15. ...
16. >>> X_train
17. array([[4, 5],
18. [0, 1],
19. [6, 7]])
20. >>> y_train
21. [2, 0, 3]
22. >>> X_test
23. array([[2, 3],
24. [8, 9]])
25. >>> y_test
26. [1, 4]

方法二

手动切分，代码如下。输入采用Python的DataFrame，同样输出三个文件。如果需要每次都输入同样的切分数据，可采用random.seed()定义随机数种子。

1. def split_train_test_valid():
2. # read file
3. input_path = "E:\\Data\\"
4. file = "flow.csv"
5. df_flow = pd.read_csv(input_path + file, header=None, encoding='gbk')
7. # define the ratios 6:2:2
8. train_len = int(len(df_flow) * 0.6)
9. test_len = int(len(df_flow) * 0.2)
11. # split the dataframe
12. idx = list(df_flow.index)
13. random.shuffle(idx) # 将index列表打乱
14. df_train = df_flow.loc[idx[:train_len]]
15. df_test = df_flow.loc[idx[train_len:train_len+test_len]]
16. df_valid = df_flow.loc[idx[train_len+test_len:]] # 剩下的就是valid
18. # output
19. df_train.to_csv(input_path+'train.txt', header=False, index=False, sep='\t')
20. df_test.to_csv(input_path+'test.txt', header=False, index=False, sep='\t')
21. df_valid.to_csv(input_path+'valid.txt', header=False, index=False, sep='\t')

参考资料：https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html