sklearn 分割训练集 验证集和测试集
时间: 2023-08-23 09:06:42 浏览: 56
在使用 sklearn 进行机器学习建模时,一般需要将数据集分为训练集、验证集和测试集三部分。训练集用于训练模型,验证集用于调参和评估模型的性能,测试集用于最终评估模型的泛化性能。sklearn 提供了多种用于分割数据集的函数和类,包括 train_test_split、StratifiedShuffleSplit、KFold 等。
下面以 train_test_split 为例,介绍如何使用 sklearn 分割训练集、验证集和测试集。
首先,导入需要的库和数据集:
```python
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
iris = load_iris()
X, y = iris.data, iris.target
```
然后,调用 train_test_split 函数进行分割,指定参数 test_size 和 random_state:
```python
X_trainval, X_test, y_trainval, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
```
这里将数据集分成了训练集 + 验证集和测试集两部分,其中训练集 + 验证集占原始数据集的 80%,测试集占 20%。random_state 参数用于设置随机种子,确保每次运行代码得到的分割结果一致。
接着,再将训练集 + 验证集分成训练集和验证集两部分:
```python
X_train, X_val, y_train, y_val = train_test_split(X_trainval, y_trainval, test_size=0.25, random_state=42)
```
这里将训练集 + 验证集再次分成了训练集和验证集,其中训练集占训练集 + 验证集的 75%,验证集占 25%。
最后,可以查看数据集的大小:
```python
print("训练集大小:", X_train.shape[0])
print("验证集大小:", X_val.shape[0])
print("测试集大小:", X_test.shape[0])
```
输出结果为:
```
训练集大小: 90
验证集大小: 30
测试集大小: 30
```
这样就完成了 sklearn 分割训练集、验证集和测试集的操作。