train_test_split資料集分割

weixin_34279579發表於2018-12-28

函式名:train_test_split
所在包:sklearn.model_selection
功能:劃分資料的訓練集與測試集
引數解讀:train_test_split (*arrays,test_size, train_size, rondom_state=None, shuffle=True, stratify=None)

  • arrays:特徵資料和標籤資料(array,list,dataframe等型別),要求所有資料長度相同。
  • test_size / train_size: 測試集/訓練集的大小,若輸入小數表示比例,若輸入整數表示資料個數。
  • rondom_state:隨機種子(一個整數),其實就是一個劃分標記,對於同一個資料集,如果rondom_state相同,則劃分結果也相同。
  • shuffle:是否打亂資料的順序,再劃分,預設True。
  • stratify:none或者array/series型別的資料,表示按這列進行分層取樣。

舉個栗子:

特徵資料:data
   a  b  c
0  1  2  3
1  1  3  6
2  2  3  8
3  1  5  7
4  2  4  8
5  2  3  6
6  1  4  8
7  2  3  6
標籤資料:label
[2,3,5,6,8,0,2,3]

#劃分
xtrain,xtest,ytrain,ytest=train_test_split(data,label,test_size=0.2,stratify=data['a'],random_state=1)
訓練特徵集:
   a  b  c
0  1  2  3
2  2  3  8
3  1  5  7
5  2  3  6
6  1  4  8
4  2  4  8
測試特徵集:
   a  b  c
1  1  3  6
7  2  3  6

訓練集與測試集按照a列來分層取樣,且無論重複多少次上述語句,劃分結果都相同。

相關文章