如何在Python中将数据集分成训练集和测试集?
在机器学习和数据挖掘领域,我们经常需要将数据集分成训练集和测试集。训练集用来训练模型,测试集用来评估模型的性能。在Python中,我们可以使用scikit-learn库中的train_test_split函数来实现这个过程。
train_test_split函数
train_test_split函数是scikit-learn库中的一个函数,可以将输入数据集随机划分为训练集和测试集。该函数的语法如下:
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0)
train_test_split函数接受4个参数:
- X:特征值,即输入的数据集。可以是一个数组或矩阵。
- y:输出的目标值,即希望模型能够预测的值。可以是数组、矩阵或Series类型。
- test_size:测试集所占的比例。默认值为0.25,即测试集占25%。
- random_state:随机数种子。在train_test_split函数中,随机数种子用于控制随机的过程,保证重复运行代码得到的结果是一样的。如果不指定此参数,则每次运行代码得到的结果都不一样。
train_test_split函数返回4个对象:
- X_train:训练集的特征值。
- X_test:测试集的特征值。
- y_train:训练集的目标值。
- y_test:测试集的目标值。
示例代码
下面是一个简单的示例代码,用于将数据集分成训练集和测试集。
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
# 加载iris数据集
iris = load_iris()
# 获取特征值和目标值
X = iris.data
y = iris.target
# 将数据集划分成训练集和测试集,测试集占30%
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0)
# 输出训练集和测试集的形状
print("训练集的形状:", X_train.shape)
print("测试集的形状:", X_test.shape)
输出结果如下:
训练集的形状: (105, 4)
测试集的形状: (45, 4)
不同应用场景下的训练集和测试集划分方法
在实际应用中,根据不同的数据集和任务要求,我们可以采用不同的训练集和测试集划分方法。下面是一些常用的划分方法。
留出法
留出法又称为简单随机划分法,即将数据集划分成两个互斥的集合:训练集和测试集。通常情况下,训练集占总样本数的70%~80%,测试集占总样本数的20%~30%。留出法适用于数据量比较大的情况,但是不能充分利用数据集,可能会导致测试结果的不稳定。
使用留出法,我们只需要调用train_test_split函数即可。下面是一个示例代码:
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
# 加载iris数据集
iris = load_iris()
# 获取特征值和目标值
X = iris.data
y = iris.target
# 将数据集划分成训练集和测试集,测试集占30%
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0)
# 输出训练集和测试集的形状
print("训练集的形状:", X_train.shape)
print("测试集的形状:", X_test.shape)
K折交叉验证法
K折交叉验证法将整个数据集划分成K个互斥的子集,每次使用其中一个子集作为测试集,其余K-1个子集作为训练集进行模型训练。重复K次,每个子集都被用于测试一次,得到K个测试结果的平均值。K折交叉验证法适用于数据量较小的情况,可以充分利用数据集,提高模型的稳定性。
K折交叉验证法可以通过调用KFold或StratifiedKFold函数实现。其中,KFold函数将数据集按顺序划分成K个子集,每个子集都作为测试集;StratifiedKFold函数保证每个子集中不同类别的样本比例与原始数据集中的比例相同。
下面是一个使用KFold函数的示例代码:
from sklearn.model_selection import KFold
from sklearn.datasets import load_iris
# 加载iris数据集
iris = load_iris()
# 获取特征值和目标值
X = iris.data
y = iris.target
# 使用5折交叉验证划分数据集
kf = KFold(n_splits=5, shuffle=True, random_state=0)
# 对每个子集进行训练和测试,并输出测试结果
for train_index, test_index in kf.split(X, y):
X_train, X_test = X[train_index], X[test_index]
y_train, y_test = y[train_index], y[test_index]
print("训练集的形状:", X_train.shape)
print("测试集的形状:", X_test.shape)
分层K折交叉验证法
分层K折交叉验证法与K折交叉验证法类似,也是将数据集划分成K个互斥的子集,每次使用其中一个子集作为测试集,其余K-1个子集作为训练集进行模型训练。但是不同之处在于,分层K折交叉验证法可以保证每个子集中不同类别的样本比例与原始数据集中的比例相同。
分层K折交叉验证法可以通过调用StratifiedKFold函数实现。下面是一个示例代码:
from sklearn.model_selection import StratifiedKFold
from sklearn.datasets import load_iris
# 加载iris数据集
iris = load_iris()
# 获取特征值和目标值
X = iris.data
y = iris.target
# 使用5折交叉验证划分数据集
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
# 对每个子集进行训练和测试,并输出测试结果
for train_index, test_index in skf.split(X, y):
X_train, X_test = X[train_index], X[test_index]
y_train, y_test = y[train_index], y[test_index]
print("训练集的形状:", X_train.shape)
print("测试集的形状:", X_test.shape)
结论
在Python中将数据集分成训练集和测试集,可以使用scikit-learn库中的train_test_split函数。根据不同的应用场景,我们可以采用留出法、K折交叉验证法或分层K折交叉验证法等不同的训练集和测试集划分方法。