分割训练集测试集验证集的两种方式

训练模型过程中，不可避免要对数据集进行切分，将数据集按比例分成训练集，测试集，验证集。作为一种常见的需求，稍微总结了一下两种实现方式。

1.train_test_split

平时最常见的切分数据集的方法是sklearn中的train_test_split方法，该方法可以很方便将数据集分成训练集与测试集。

import random
import pandas as pd
from sklearn.model_selection import train_test_split


def gen_data():
    x1 = [random.random() for _ in range(100)]
    x2 = [random.randint(0, 9) for _ in range(100)]
    y = [random.randint(0, 1) for _ in range(100)]
    data = pd.DataFrame({'f1': x1, 'f2': x2, 'label': y})

    X_train, X_test, y_train, y_test = train_test_split(data[['f1', 'f2']], data.label, test_size=0.2)
    print(X_train)
    print()
    print(X_test)
    print()
    print(y_train)
    print()
    print(y_test)

上面的代码运行以后，结果为

          f1  f2
58  0.167506   4
75  0.687526   8
26  0.973927   9
5   0.959837   9
28  0.713963   0
..       ...  ..
42  0.513807   2
16  0.398464   6
90  0.809922   4
3   0.957735   2
47  0.822975   3

[80 rows x 2 columns]

          f1  f2
59  0.148613   7
40  0.981047   8
6   0.609841   3
7   0.364588   5
24  0.907005   8
92  0.530647   8
99  0.675087   0
37  0.131263   7
63  0.187509   2
39  0.657692   1
56  0.932423   0
23  0.800484   8
54  0.428490   5
25  0.853351   6
17  0.525853   7
64  0.045794   4
4   0.642669   7
50  0.149979   4
87  0.969403   6
77  0.096266   9

58    0
75    1
26    0
5     1
28    0
     ..
42    1
16    1
90    1
3     0
47    1
Name: label, Length: 80, dtype: int64

59    0
40    1
6     0
7     1
24    1
92    0
99    0
37    0
63    0
39    0
56    1
23    0
54    0
25    0
17    0
64    1
4     0
50    0
87    0
77    1
Name: label, dtype: int64

2.自己切分实现

train_test_split确实很方便，也很实用。但是该方法有一个问题，只能切分成train,test两个数据集。很多时候我们还需要一个验证集，这个时候train_test_split方法就无能为力了。因此我们自己简单实现一个切分的方法。

import random
import pandas as pd

def gen_train_test_val():
    x1 = [random.random() for _ in range(100)]
    x2 = [random.randint(0, 9) for _ in range(100)]
    y = [random.randint(0, 1) for _ in range(100)]
    data = pd.DataFrame({'f1': x1, 'f2': x2, 'label': y})

    train_size = int(len(data) * 0.7)
    test_size = int(len(data) * 0.2)
    idx = list(data.index)
    random.shuffle(idx)

    df_train = data.loc[idx[:train_size]]
    df_test = data.loc[idx[train_size:train_size+test_size]]
    df_valid = data.loc[idx[train_size+test_size:]]

    print(df_train)
    print()
    print(df_test)
    print()
    print(df_valid)
    print()


gen_train_test_val()

上面实现方式的核心思路是，利用dataframe默认的index，做shuffle。shuffle完毕以后，根据设定的train_size，test_size来选取数据集。比如在上面的实现中，我们就是设定训练集70%，测试集20%，验证集为剩余的10%。

最后代码输出为

          f1  f2  label
64  0.397632   5      1
99  0.938272   3      1
26  0.832233   5      1
48  0.259652   6      0
57  0.083901   8      1
..       ...  ..    ...
9   0.629019   1      0
81  0.261926   4      0
54  0.691443   5      1
86  0.617717   7      1
45  0.357114   5      0

[70 rows x 3 columns]

          f1  f2  label
52  0.693583   2      1
12  0.431800   9      0
5   0.865485   5      0
96  0.273552   8      1
58  0.739783   7      1
87  0.245566   4      1
65  0.468828   3      1
94  0.782948   7      0
69  0.393221   1      0
17  0.388766   6      1
23  0.829479   9      1
8   0.052652   7      0
93  0.458537   3      1
55  0.174687   4      1
83  0.857251   3      0
14  0.547664   7      1
13  0.037768   4      0
41  0.626701   4      0
50  0.289869   5      0
75  0.889851   3      1

          f1  f2  label
51  0.834317   3      0
70  0.885837   4      0
74  0.343928   1      0
44  0.707159   1      1
85  0.108250   3      0
88  0.045994   7      0
53  0.472742   9      0
24  0.498392   9      0
82  0.856740   1      1
16  0.608382   7      1