如何用scikit-learn的train_test_split()函数分割数据集

在这篇文章中，我们将讨论如何用scikit-learns的train_test_split()分割数据集。

sklearn.model_selection.train_test_split（）函数

train_test_split()方法是用来将我们的数据分成训练集和测试集。首先，我们需要将我们的数据分为特征（X）和标签（Y）。数据帧被分为X_train、X_test、y_train和y_test。X_train和y_train集被用于训练和拟合模型。X_test和y_test集用于测试模型是否预测了正确的输出/标签。我们可以明确地测试训练和测试集的大小。建议保持我们的训练集比测试集大。

• Train set: 训练数据集是一组被用来拟合模型的数据。训练模型的数据集。这些数据被模型看到并学习。
• Test set: 测试数据集是训练数据集的一个子集，利用它来对最终的模型拟合作出准确的评价。
• validation set:验证数据集是来自模型训练集的数据样本，用于估计模型的性能，同时调整模型的超参数。
• underfitting: 拟合不足的数据模型在训练集和未观察到的数据上都有很高的错误率，因为它不能有效地表示输入和输出变量之间的关系。
• overfitting: 当一个统计模型与它的训练数据完全匹配，但由于算法无法针对未见过的数据准确执行而失去了目标，这被称为过度拟合

语法: sklearn.model_selection.train_test_split(*arrays, test_size=None, train_size=None, random_state=None, shuffle=True, stratify=None

参数:

• arrays：可索引的序列。 Lists, numpy arrays, scipy-sparse matrices, and pandas dataframes都是有效输入。
• test_size：int 或 float，默认为 None。 如果是float，它应该在0.0和1.0之间，代表要测试分割的数据集的百分比。如果使用int，它指的是测试样本的总数。如果该值为None，则使用训练规模的补码。如果训练规模也是无，它将被设置为0.25。
• train_size：int或float，默认为None。
• random_state : int，默认为None。 控制在执行分割之前如何洗刷数据。如果要在多个函数调用中重复输出，请传递一个int。
• shuffle: boolean对象，默认为True。 分割前是否要对数据进行洗牌。如果shuffle=False，Stratify必须为None。
• stratify：数组类对象，默认为None。如果选择None，则使用这些作为类标签对数据进行分层。

返回值:

splitting: 输入的训练-测试分割被表示为一个列表。

拆分数据集的步骤

第一步：导入必要的包或模块

在这一步中，我们要将必要的包或模块导入工作的python环境中。

# import packages
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

第2步：导入数据框架/数据集

这里，我们使用pd.read_csv()方法从pandas加载CSV，并使用shape()函数获得数据集的形状。

CSV Used:

# importing data
df = pd.read_csv('prediction.csv')
print(df.shape)

输出:

(13, 3)

第三步：获取X和Y特征变量

这里，我们要分配X和Y变量，其中X特征变量有自变量，Y特征变量有因变量。

X= df['area']
y=df['prices']

第四步：使用train test split类将数据分成训练集和测试集

这里，来自sklearn.model_selection的train_test_split()类被用来将我们的数据分成训练集和测试集，其中特征变量被作为方法中的输入。test_size决定了将进入测试集的数据部分，随机状态被用于数据的可重复性。

# using the train test split function
X_train, X_test, y_train, y_test = train_test_split(
  X,y , random_state=104,test_size=0.25, shuffle=True)

示例:

在这个例子中，’predictions.csv’文件被导入。df.shape属性被用来检索数据框的形状。数据框的形状是（13,3）。特征列在X变量中，结果列在y变量中。X和y变量被传递到train_test_split()方法中，以将数据帧分成训练集和测试集。随机状态参数用于数据的可重复性。test_size给定为0.25，这意味着25%的数据会进入测试集。数据框中的13行中有4行进入测试集。75%的数据进入训练集，也就是13行中的9行。训练集被用来适应和训练机器学习模型。测试集则用于评估。

CSV Used:

# import packages
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
  
# importing data
df = pd.read_csv('prediction.csv')
print(df.shape)
  
# head of the data
print('Head of the dataframe : ')
print(df.head())
  
print(df.columns)
  
X= df['area']
y=df['prices']
  
# using the train test split function
X_train, X_test, y_train, y_test = train_test_split(
  X,y , random_state=104,test_size=0.25, shuffle=True)
  
# printing out train and test sets
  
print('X_train : ')
print(X_train.head())
print(X_train.shape)
  
print('')
print('X_test : ')
print(X_test.head())
print(X_test.shape)
  
print('')
print('y_train : ')
print(y_train.head())
print(y_train.shape)
  
print('')
print('y_test : ')
print(y_test.head())
print(y_test.shape)

输出:

(13, 3)
Head of the dataframe : 
   Unnamed: 0  area         prices
0           0  1000  316404.109589
1           1  1500  384297.945205
2           2  2300  492928.082192
3           3  3540  661304.794521
4           4  4120  740061.643836
Index(['Unnamed: 0', 'area', 'prices'], dtype='object')
X_train : 
3    3540
7    3460
4    4120
0    1000
8    4750
Name: area, dtype: int64
(9,)

X_test : 
12    7100
2     2300
11    8600
10    9000
Name: area, dtype: int64
(4,)

y_train : 
3    661304.794521
7    650441.780822
4    740061.643836
0    316404.109589
8    825607.876712
Name: prices, dtype: float64
(9,)

y_test : 
12    1.144709e+06
2     4.929281e+05
11    1.348390e+06
10    1.402705e+06
Name: prices, dtype: float64
(4,)

示例:

在这个例子中，执行了以下步骤：

• 导入了必要的软件包。
• 广告.csv数据集被加载和清理，空值被删除。
• 创建特征和目标数组（X和y）。
• 创建的数组被分成训练集和测试集。30%的数据集被放入测试集，这意味着70%的数据是训练集。
• 创建一个标准的缩放器对象。
• X_train被装入标度器。
• 使用transform()方法对X_train和X_test进行转换。
• 一个简单的线性回归模型被创建
• Train集适合于模型。
• 使用predict()方法来对X_test集进行预测。
• mean_squared_error()指标被用来评估模型。

要查看和下载本例中使用的CSV文件，请点击这里。

# import packages
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
  
df = pd.read_csv('Advertising.csv')
# dropping rows which have null values
df.dropna(inplace=True,axis=0)
  
y = df['sales']
X = df.drop('sales',axis=1)
  
# splitting the dataframe into train and test sets
X_train,X_test,y_train,y_test = train_test_split(
  X,y,test_size=0.3,random_state=101)
scaler = StandardScaler()
scaler.fit(X_train)
X_train = scaler.transform(X_train)
X_test = scaler.transform(X_test)
  
model = LinearRegression().fit(X_train,y_train)
y_pred = model.predict(X_test)
print(y_pred)
print(mean_squared_error(y_test,y_pred))

输出:

array([19.82000933, 14.23636718, 12.80417236, 7.75461569, 8.31672266,

15.4001915 , 11.6590983 , 15.22650923, 15.53524916, 19.46415132,

17.21364106, 16.69603229, 16.46449309, 10.15345178, 13.44695953,

24. 71946196, 18.67190453, 15.85505154, 14.45450049, 9.91684409,

10.41647177, 4.61335238, 17.41531451, 17.31014955, 21. 72288151,

5.87934089, 11.29101265, 17.88733657, 21.04225992, 12.32251227,

14.4099317 , 15.05829814, 10. 2105313 , 7.28532072, 12.66133397,

23.25847491, 18.87101505, 4.55545854, 19.79603707, 9.21203026,

10.24668718, 8.96989469, 13.33515217, 20.69532628, 12.17013119,

21.69572633, 16.7346457 , 22.16358256, 5. 34163764, 20.43470231,

7.58252563, 23.38775769, 10.2270323 , 12.33473902, 24.10480458,

9.88919804, 21.7781076 ])

2.7506859249500466

示例:

在这个例子中，我们要使用K-近邻分类器模型。

在这个例子中，执行了以下步骤：

• 必要的包被导入。
• 从sklearn.datasets加载虹膜数据。
• 特征和目标数组被创建（X和y）。
• 创建的数组被分成训练集和测试集。30%的数据集被放入测试集，这意味着70%的数据是训练集。
• 使用KNeighborsClassifier类创建一个基本的Knn模型。
• 训练集适合在Knn模型中。
• predict()方法被用来对X_test集进行预测。

# Import packages
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
   
# Load the data
irisData = load_iris()
   
# Create feature and target arrays
X = irisData.data
y = irisData.target
   
# Split data into train and test sets
X_train, X_test, y_train, y_test = train_test_split(
             X, y, test_size = 0.2, random_state=42)
   
knn = KNeighborsClassifier(n_neighbors=1)
   
knn.fit(X_train, y_train)
   
# predicting on the X_test data set
print(knn.predict(X_test))

输出:

[1 0 2 1 1 0 1 2 1 1 2 0 0 0 0 1 2 1 1 2 0 2 0 2 2 2 2 2 0 0]