Python 如何拆分数据集

前言

训练模型之前一般需要把数据集拆分为训练集和测试集，使用python代码如何拆分的关键就是如何更方便的选择出自变量X和因变量Y。

加载数据

# 导入第三方模块
import pandas as pd
# 读入数据
df = pd.read_csv(r'splitfeatures.csv')
df.head()

如何选择出X和Y

1. 最简单的方式一个一个选择

X = df[['age','sex','income','profession']]
y = df['is_buy']

2. 用iloc和loc选择

X1 = df.iloc[:,[0,1,2,4]]
y1 = df.iloc[:,[3]]
X2 = df.loc[:,['age','sex','income','profession']]
y2 = df.loc[:,['is_buy']]

3. 用drop选择

X3 = df.drop(['is_buy'],axis=1)
y3 = df.is_buy

4. 使用dataframe的colunms方法

col = df.columns[[0,1,2,4]]
X4 = df.loc[:,col]
X5 = df.iloc[:, df.columns != 'is_buy']
y5 = df.iloc[:, df.columns == 'is_buy']

使用model_selection拆分数据集

from sklearn import model_selection
# 将数据集拆分为训练集和测试集
X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y, test_size = 0.2, random_state = 1234)