机器学习笔记——数据集分割

在模型训练之前，要首先划分训练集与测试集，如何对原始数据集进行训练集与测试集的划分？训练集与测试集的比例各占多少？如何保证各自内部标签分布平衡都会影响模型训练的最终效果。

好在R和Python中有现成的数据集分割函数，避免手动写函数导致划分比例不合理、训练集与测试集的样本的结构与总体不均衡的问题。

R语言中caTools包中的sample.split函数可以用来自动将原始数据集分割成训练集和测试集。

方法一——caTools中的sample.split函数

library("caTools") set.seed(123) data(iris) table(iris$Species) setosa versicolor virginica 50 50 50

split = sample.split(iris$Species,SplitRatio = .8) train_data = subset(iris,split == TRUE) test_data = subset(iris,split == FALSE) table(train_data$Species) setosa versicolor virginica 40 40 40

table(test_data$Species) setosa versicolor virginica 10 10 10

划分方法二——createDataPartition函数

library("caret")
split1 <- createDataPartition(y=iris$Species,p=0.8,list = FALSE)
 
train_data <- iris[split1,]
table(train_data$Species)
 setosa versicolor virginica 
 40 40 40 
test_data <- iris[-split1,]
table(test_data$Species)
 setosa versicolor virginica 
 10 10 10

可以看到无论是caTools包中的sample.split函数还是caret包中的createDataPartition函数，都针对分类标签做了混合后的分层随机抽样，这样可以保证训练集与测试集内的各类标签分布比例与样本总体的分布比例严格一致，否则如果仅仅使用sample函数，无法达到分层随机抽样的目的。

Python的sk-learn库中也有现成的数据集分割工具可用。

from sklearn.model_selection import train_test_splitfrom sklearn.datasets import load_irisimport pandas as pd
 
iris = load_iris()
data = iris['data']
iris_data = pd.DataFrame(
 data = data,
 columns = ['sepal_length','sepal_width','petal_length','petal_width']
 )
iris_data["Species"] = iris[ 'target']
iris_data["Species"] = iris_data["Species"].map({0:"setosa",1:"versicolor",2:"virginica"})
 
iris_data["Species"].value_counts()
versicolor 50virginica 50setosa 50Name: Species, dtype: int64
 
x,y = iris_data.iloc[:,0:-1],iris_data.iloc[:,-1]
train_data,test_data,train_target,test_target = train_test_split(x,y,test_size = 0.2,stratify = y)
 
train_target.value_counts()
test_target.value_counts()

这里有一个细节，其中train_test_split函数中有两个特殊的参数：shuffle和stratify。

shuffle参数相当于对原始数据进行混合抽样，相当于扑克牌发牌前的洗牌，保证随机性。

stratify参数则可以保证训练集&测试集中样本标签结构比例与指定的总体中样本标签结构比例一致，特别是在原始数据中样本标签分布不均衡时非常有用，达到分层随机抽样的目的。

而对于训练集与测试集的划分比例上来看，目前没有一个统一的比例标准，在机器学习类的参考资料中，推荐的比例是训练集和测试集的比例保持在1:2~1:4之间为佳。

倘若使用验证集，则建议保持6:2:2左右的训练集、验证集及测试集比例。

原文发布时间为：2018-09-23

本文作者：杜雨

本文来自云栖社区合作伙伴“数据小魔方”，了解相关信息可以关注“数据小魔方”。