sklearn.model_selection.train_test_split

左手的ㄟ右手 2022-05-21 06:36 55阅读 0赞

数据集划分：sklearn.model\_selection.train\_test\_split(\*arrays, \*\*options)

主要参数说明：

\*arrays：可以是列表、numpy数组、scipy稀疏矩阵或pandas的数据框

test\_size：可以为浮点、整数或None，默认为None

①若为浮点时，表示测试集占总样本的百分比

②若为整数时，表示测试样本样本数

③若为None时，test size自动设置成0.25

train\_size：可以为浮点、整数或None，默认为None

①若为浮点时，表示训练集占总样本的百分比

②若为整数时，表示训练样本的样本数

③若为None时，train\_size自动被设置成0.75

random\_state：可以为整数、RandomState实例或None，默认为None

①若为None时，每次生成的数据都是随机，可能不一样

②若为整数时，每次生成的数据都相同

stratify：可以为类似数组或None

①若为None时，划分出来的测试集或训练集中，其类标签的比例也是随机的

②若不为None时，划分出来的测试集或训练集中，其类标签的比例同输入的数组中类标签的比例相同，可以用于处理不均衡的数据集

通过简单栗子看看各个参数的作用：

①test\_size决定划分测试、训练集比例

1.  In \[ 1\]: import numpy as np
2.  ...: from sklearn.model\_selection import train\_test\_split
3.  ...: X = np.arange( 20)
4.  ...: y = \[ 'A', 'B', 'A', 'A', 'A', 'B', 'A', 'B', 'B', 'A', 'A', 'B', 'B', 'A', 'A', 'B', 'A
5.   ...: ', 'B', 'A', 'A'\]
6.  ...: X\_train , X\_test , y\_train,y\_test = train\_test\_split(X,y,test\_size= 0.25
7.  ...: ,random\_state= 0)
8.  ...:
9.  
10. In \[ 2\]: X\_test.shape
11. Out\[ 2\]: ( 5,)
12. 
13. In \[ 3\]: X\_train.shape
14. Out\[ 3\]: ( 15,)
15. 
16. In \[ 4\]: X\_test ,y\_test
17. Out\[ 4\]: (array(\[ 18, 1, 19, 8, 10\]), \[ 'A', 'B', 'A', 'B', 'A'\])

②random\_state不同值获取到不同的数据集

设置random\_state=0再运行一次，结果同上述相同

1.  In \[ 5\]: import numpy as np
2.  ...: from sklearn.model\_selection import train\_test\_split
3.  ...: X = np.arange( 20)
4.  ...: y = \[ 'A', 'B', 'A', 'A', 'A', 'B', 'A', 'B', 'B', 'A', 'A', 'B', 'B', 'A', 'A', 'B', 'A
5.   ...: ', 'B', 'A', 'A'\]
6.  ...: X\_train , X\_test , y\_train,y\_test = train\_test\_split(X,y,test\_size= 0.25
7.  ...: ,random\_state= 0)
8.  ...: X\_test ,y\_test
9.  ...:
10. Out\[ 5\]: (array(\[ 18, 1, 19, 8, 10\]), \[ 'A', 'B', 'A', 'B', 'A'\])

设置random\_state=None运行两次，发现两次的结果不同

1.  In \[ 6\]: import numpy as np
2.  ...: from sklearn.model\_selection import train\_test\_split
3.  ...: X = np.arange( 20)
4.  ...: y = \[ 'A', 'B', 'A', 'A', 'A', 'B', 'A', 'B', 'B', 'A', 'A', 'B', 'B', 'A', 'A', 'B', 'A
5.   ...: ', 'B', 'A', 'A'\]
6.  ...: X\_train , X\_test , y\_train,y\_test = train\_test\_split(X,y,test\_size= 0.25
7.  ...: )
8.  ...: X\_test ,y\_test
9.  ...:
10. Out\[ 6\]: (array(\[ 3, 18, 14, 7, 4\]), \[ 'A', 'A', 'A', 'B', 'A'\])
11. 
12. In \[ 7\]: import numpy as np
13. ...: from sklearn.model\_selection import train\_test\_split
14. ...: X = np.arange( 20)
15. ...: y = \[ 'A', 'B', 'A', 'A', 'A', 'B', 'A', 'B', 'B', 'A', 'A', 'B', 'B', 'A', 'A', 'B', 'A
16.  ...: ', 'B', 'A', 'A'\]
17. ...: X\_train , X\_test , y\_train,y\_test = train\_test\_split(X,y,test\_size= 0.25
18. ...: )
19. ...: X\_test ,y\_test
20. ...:
21. Out\[ 7\]: (array(\[ 18, 6, 3, 14, 8\]), \[ 'A', 'A', 'A', 'A', 'B'\])

③设置stratify参数，可以处理数据不平衡问题

1.  In \[ 8\]: import numpy as np
2.  ...: from sklearn.model\_selection import train\_test\_split
3.  ...: X = np.arange( 20)
4.  ...: y = \[ 'A', 'B', 'A', 'A', 'A', 'B', 'A', 'B', 'B', 'A', 'A', 'B', 'B', 'A', 'A', 'B', 'A
5.   ...: ', 'B', 'A', 'A'\]
6.  ...: X\_train , X\_test , y\_train,y\_test = train\_test\_split(X,y,test\_size= 0.25
7.  ...: ,stratify=y)
8.  ...: X\_test ,y\_test
9.  ...:
10. Out\[ 8\]: (array(\[ 18, 8, 3, 10, 11\]), \[ 'A', 'B', 'A', 'A', 'B'\])
11. 
12. In \[ 9\]: import numpy as np
13. ...: from sklearn.model\_selection import train\_test\_split
14. ...: X = np.arange( 20)
15. ...: y = \[ 'A', 'B', 'A', 'A', 'A', 'B', 'A', 'B', 'B', 'A', 'A', 'B', 'B', 'A', 'A', 'B', 'A
16.  ...: ', 'B', 'A', 'A'\]
17. ...: X\_train , X\_test , y\_train,y\_test = train\_test\_split(X,y,test\_size= 0.25
18. ...: ,stratify=y)
19. ...: X\_test ,y\_test
20. ...:
21. Out\[ 9\]: (array(\[ 6, 19, 8, 17, 0\]), \[ 'A', 'A', 'B', 'B', 'A'\])
22. 
23. In \[ 10\]: X\_train,y\_train
24. Out\[ 10\]:
25. (array(\[ 7, 1, 11, 10, 15, 2, 3, 5, 4, 13, 12, 16, 18, 14, 9\]),
26. \[ 'B', 'B', 'B', 'A', 'B', 'A', 'A', 'B', 'A', 'A', 'B', 'A', 'A', 'A', 'A'\])

设置stratify=y时，我们发现每次划分后，测试集和训练集中的类标签比例同原始的样本中类标签的比例相同，都为2:3