文章目录

• train_test_split()用法
• 获取数据
• 划分训练集和测试集
• 完整代码脚手架

train_test_split()用法

python机器学习中常用 train_test_split()函数划分训练集和测试集，其用法语法如下：

• X_train, X_test, y_train, y_test = train_test_split(train_data, train_target, test_size, random_state, shuffle)

以sklearn库内置的iris数据集（鸢尾数据集）为例，首先获取数据：

获取数据

from sklearn.model_selection import train_test_split
# 以sklearn库内置的iris数据集（鸢尾数据集）为例
dataset = load_iris()

这里的dataset数据是sklearn.utils.Bunch类型的数据，比较像字典

将其打印出~

print(dataset)

如下所示

从中取出其data属性和target属性，X是特征数组（也称数据集），y表示类别数组（也称标签）

X = dataset.data
y = dataset.target

此例中，有四个特征（即data的4列表示4个特征），分别是鸢尾植物的萼片的长，萼片的宽，花瓣的长，花瓣的宽。
X中共150行，即150个样本，类别数据总共有150个数据(对应150个样本的类别)。

print(y)

y的150个数据如上图，其中，有0,1,2三个取值，表示三种花：

使用最简单的离散化算法，以均值为阈值，使大于阈值的特征值为1，小于阈值的特征值为0.

attribute_means = X.mean(axis=0) 
X_d = np.array(X >= attribute_means, dtype='int')  # bool转换
print(X_d) 

运行结果（成功将X的数据转换为bool类型）：

划分训练集和测试集

然后就是使用train_test_split()函数将数据划分训练集和测试集了。

random_state = 10
# 以默认的0.25作为分割比例进行分割（训练集:测试集=3:1）
X_train, X_test, y_train, y_test = train_test_split(X_d, y, random_state=random_state)
print("There are {} training samples".format(y_train.shape[0]))
print("There are {} testing samples".format(y_test.shape[0]))

如图得到的数据中112/38接近3:1。分割成功！

完整代码脚手架

（将上述分步的代码写在一块儿方便复制使用）：

from sklearn.model_selection import train_test_split
dataset = load_iris()

X = dataset.data
y = dataset.target

attribute_means = X.mean(axis=0) 
X_d = np.array(X >= attribute_means, dtype='int')

random_state = 10
X_train, X_test, y_train, y_test = train_test_split(X_d, y, random_state=random_state)