python train_test_split stratify
时间: 2023-11-08 21:05:05 浏览: 126
数据集分割train和test程序
`train_test_split` 是一个用于将数据集划分为训练集和测试集的函数,`stratify` 参数是用于指定按照某个变量进行分层抽样的参数。
具体来说,如果我们希望在划分训练集和测试集时,保持原始数据中某个变量的分布比例不变,就可以使用 `stratify` 参数。例如,在分类问题中,我们希望训练集和测试集中各类别样本的比例与原始数据集中各类别样本的比例相同,就可以使用 `stratify` 参数。
下面是一个使用 `train_test_split` 函数进行分层抽样的示例代码:
```python
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y)
```
其中,`X` 是特征矩阵,`y` 是目标变量。`test_size` 参数指定测试集占总样本数的比例,`stratify` 参数指定按照哪个变量进行分层抽样。
阅读全文