機器學習 train_test_split參數含義

woff

在機器學習中，我們通常將原始數據按照比例分割為“測試集”和“訓練集”，通常使用sklearn.cross_validation裡的train_test_split模塊用來分割數據。

cross_validation已經棄用，現在改為從sklearn.model_selection中調用train_test_split函數。

簡單用法如下：

X_train,X_test, y_train, y_test =sklearn.model_selection.train_test_split(train_data,train_target,test_size=0.4, random_state=0, stratify= y_train )
# train_data：所要劃分的樣本特徵集
# train_target：所要劃分的樣本結果
# test_size：樣本佔比，如果是整數的話就是樣本的數量
# random_state：是隨機數的種子。
# 隨機數種子：其實就是該組隨機數的編號，在需要重複試驗的時候，保證得到一組一樣的隨機數。比如你每次都填1，其他參數一樣的情況下你得到的隨機數組是一樣的。但填0或不填，每次都會不一樣。

stratify是為了保持split前類的分佈。比如有100個數據，80個屬於A類，20個屬於B類。如果train_test_split(... test_size=0.25, stratify = y_all),那麼split之後數據如下：
training: 75個數據，其中60個屬於A類，15個屬於B類。
testing: 25個數據，其中20個屬於A類，5個屬於B類。

用了stratify參數，training集和testing集的類的比例是A：B= 4：1，等同於split前的比例（80：20）。通常在這種類分佈不平衡的情況下會用到stratify。

將stratify=X就是按照X中的比例分配

將stratify=y就是按照y中的比例分配


文章出處