train_test_split,用一个具体的,生动形象的案例给我举例一下。
时间: 2023-07-20 08:10:55 浏览: 43
假设我们有一个数据集,包含1000个人的身高和体重信息,我们想要根据这些信息来预测一个人是否为肥胖。我们将身高和体重作为特征,肥胖与否作为目标变量。
首先,我们需要将数据集划分为训练集和测试集。为了保证模型的泛化能力,我们通常会采用随机划分的方法。假设我们将数据集按照7:3的比例划分为训练集和测试集,即训练集包含700个样本,测试集包含300个样本。
我们可以使用`train_test_split()`函数来完成数据集的划分,代码如下:
```python
from sklearn.model_selection import train_test_split
X = data[['height', 'weight']]
y = data['obesity']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0)
```
其中,`X`是特征矩阵,`y`是目标变量。`test_size`参数指定了测试集占总样本数的比例,`random_state`参数指定了随机数种子,以确保每次划分的结果相同。
划分完成后,我们就可以使用训练集来训练模型,例如使用逻辑回归模型:
```python
from sklearn.linear_model import LogisticRegression
clf = LogisticRegression()
clf.fit(X_train, y_train)
```
训练完成后,我们可以使用测试集来评估模型的性能:
```python
y_pred = clf.predict(X_test)
accuracy = clf.score(X_test, y_test)
```
其中,`y_pred`是模型对测试集的预测结果,`accuracy`是模型在测试集上的准确率。我们可以根据准确率来选择最优的模型。