sklearn训练集验证集测试集划分

2,369 阅读1分钟

因为做实验需要手动处理一些原始的数据集,需要涉及到数据集的划分,就查了一些资料,发现sklearn提供有两种划分函数,一种随机划分,一种分层划分。

所谓随机划分,就是说完全随机的,不能保证训练集和测试集的不同类别的比例保持一致。而分层划分能做到这一点,能更好地体现独立同分布。

这本身没什么技术含量。但是希望本文能节省大家的时间,拷贝小修就能用,从而将更多的时间花在解决更重要的科研问题上面。

将分层划分的示例代码放在下面:

假设有两个ndarray: X, y

from sklearn.model_selection import StratifiedShuffleSplit

train_list, val_list, test_list = None, None, None

split = StratifiedShuffleSplit(n_splits=1, test_size=0.1, random_state=42)

for train_index, test_index in split.split(X, y):
    train_list = train_index
    test_list = test_index
    
split = StratifiedShuffleSplit(n_splits=1, test_size=2. / 9, random_state=42)

for train_index, val_index in split.split(X[train_list], y[train_list]):
    train_tmp = train_list[train_index]
    val_list = train_list[val_index]
    train_list = train_tmp
    
# now train_list, val_list, test_list contains the indices of train, val and test set respectively

n_splits: 表示下面的那个for循环需要多少次迭代,即进行多少次划分

test_size: 测试集占多少比例

random_state: 随机种子

上面的代码是按照7:2:1的比例进行的划分。