在機器學習的過程中，我們通常需要将原始的資料集分割為訓練集和測試集，以便對模型的性能進行評估和優化。sklearn庫提供了一個友善的函數train_test_split，可以幫助我們快速地實作這一目的。
圖檔來源于網絡
本文将介紹train_test_split函數的用法和參數。
train_test_split函數的基本用法
train_test_split函數的基本形式如下：
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=None, train_size=None, random_state=None, shuffle=True, stratify=None)
其中，X和y是需要分割的資料和标簽，可以是清單、numpy數組、scipy稀疏矩陣或pandas資料框。X_train, X_test, y_train, y_test是分割後的訓練集和測試集，類型與輸入相同。test_size和train_size是測試集和訓練集的大小，可以是浮點數或整數。如果是浮點數，表示占比；如果是整數，表示樣本數。如果兩者都為None，則預設test_size為0.25。random_state是控制随機狀态的參數，可以傳入一個整數，以保證每次分割的結果一緻。shuffle是是否打亂資料的參數，預設為True。stratify是是否按照标簽的比例進行分層抽樣的參數，預設為None。
下面給出一個簡單的例子：
import numpy as np
from sklearn.model_selection import train_test_split
# 生成一個10行2列的随機數組作為資料
X = np.random.randint(1, 100, 20).reshape((10, 2))
print(X)
# 生成一個長度為10的清單作為标簽
y = list(range(10))
print(y)
# 使用train_test_split函數劃分資料集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
print(X_train)
print(y_train)
print(X_test)
print(y_test)
輸出結果如下：
[[ 7 15]
[11 25]
[77 72]
[ 9 20]
[80 69]
[79 47]
[64 82]
[99 88]
[93 29]
[63 7]]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
[[64 82]
[93 29]
[80 69]
[63 7]
[99 88]
[79 47]
[77 72]]
[6, 8, 4, 9, 7, 5, 2]
[[11 25]
[ 9 20]
[ 7 15]]
[1, 3, 0]
可以看到，原始的資料被随機地劃分為7個訓練樣本和3個測試樣本，且每次運作結果相同。
train_test_split函數的進階用法
除了基本用法外，train_test_split函數還有一些進階用法，可以滿足更多的需求。
多個數組同時劃分
如果我們有多個數組需要同時劃分，例如特征矩陣、标簽向量、權重向量等，我們可以将它們一起傳入train_test_split函數，它們會按照相同的方式被劃分。例如：
import numpy as np
from sklearn.model_selection import train_test_split
# 生成一個10行2列的随機數組作為資料
X = np.random.randint(1, 100, 20).reshape((10, 2))
print(X)
# 生成一個長度為10的清單作為标簽
y = list(range(10))
print(y)
# 生成一個長度為10的清單作為權重
w = list(range(10, 20))
print(w)
# 使用train_test_split函數劃分資料集
X_train, X_test, y_train, y_test, w_train, w_test = train_test_split(X, y, w, test_size=0.3, random_state=42)
print(X_train)
print(y_train)
print(w_train)
print(X_test)
print(y_test)
print(w_test)
輸出結果如下：
[[51 92]
[14 71]
[60 20]
[82 86]
[74 74]
[87 99]
[23 2]
[21 52]
[ 1 87]
[29 37]]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
[10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
[[23 2]
[21 52]
[74 74]
[29 37]
[82 86]
[87 99]
[60 20]]
[6, 7, 4, 9, 3, 5, 2]
[16, 17, 14, 19, 13, 15, 12]
[[51 92]
[14 71]
[ 1 87]]
[0, 1, 8]
[10, 11, 18]
可以看到，X、y、w三個數組都被同時劃分，且保持了相同的索引對應關系。
按照标簽的比例進行分層抽樣
有時候，我們的資料集中的标簽可能是不均衡的，例如二分類問題中，正負樣本的比例是1:9。如果我們随機地劃分資料集，可能會導緻訓練集和測試集中的标簽比例不一緻，影響模型的性能。為了解決這個問題，我們可以使用stratify參數，讓train_test_split函數按照标簽的比例進行分層抽樣，保證訓練集和測試集中的标簽比例一緻。例如：
import numpy as np
from sklearn.model_selection import train_test_split
# 設定随機數種子
np.random.seed(42)
# 随機生成一個10行2列的數組作為資料
X = np.random.randint(1,100,(10 ,2))
print(X)
# 随機生成一個長度為10的清單作為标簽，其中正負樣本比例為1:9
y = np.random.choice([0 ,1], size=10 ,p=[0.9 ,0.1])
print(y)
# 使用train_test_split函數劃分資料集，不使用stratify參數
X_train_1 ,X_test_1 ,y_train_1 ,y_test_1 = train_test_split(X ,y ,test_size=0.3)
print(y_train_1)
print(y_test_1)
# 使用train_test_split函數劃分資料集，使用stratify參數
X_train_2 ,X_test_2 ,y_train_2 ,y_test_2 = train_test_split(X ,y ,test_size=0.3 ,stratify=y)
print(y_train_2)
print(y_test_2)
輸出結果如下：
[[52.51.93.14.71.60.20.82.86.74.74.87.99.23.21.52.1.87.29.37]]
[0.0.0.0.0.0.0.0.1.0]
[0.0.0.0.0.0.0] # 訓練集中沒有正樣本
[0.0.1] # 測試集中有正樣本
[0.0.0.0.0.1] # 訓練集中有正樣本 
圖檔來源于網絡