scikit-learn中的分層次訓練/測試分裂

Question

我需要將我的數據分成訓練集（75％）和測試集（25％）。我目前這樣做與下面的代碼：

X, Xt, userInfo, userInfo_train = sklearn.cross_validation.train_test_split(X, userInfo) 

但是，我想分層我的訓練數據集。我怎麼做？我一直在研究StratifiedKFold方法，但不會讓我指定75％/ 25％的分割，並且只對訓練數據集進行分層。

Answer 1

[更新爲0.17]

見sklearn.model_selection.train_test_split文檔：

from sklearn.model_selection import train_test_split 
X_train, X_test, y_train, y_test = train_test_split(X, y, 
                stratify=y, 
                test_size=0.25) 

[0.17 /更新]

有一個拉請求here。 但你可以簡單地做train, test = next(iter(StratifiedKFold(...))) 並使用火車和測試指數，如果你想。

Answer 2

TL; DR：使用StratifiedShuffleSplit與test_size=0.25

Scikit學習提供了兩個模塊分層劈裂：

1. StratifiedKFold：此模塊是作爲直接k折交叉驗證操作者有用：如其中將設置n_folds培訓/測試集，以使班級在兩者中平等。

赫雷什一些代碼（從上面的文檔直接）

>>> skf = cross_validation.StratifiedKFold(y, n_folds=2) #2-fold cross validation 
>>> len(skf) 
2 
>>> for train_index, test_index in skf: 
... print("TRAIN:", train_index, "TEST:", test_index) 
... X_train, X_test = X[train_index], X[test_index] 
... y_train, y_test = y[train_index], y[test_index] 
... #fit and predict with X_train/test. Use accuracy metrics to check validation performance 

StratifiedShuffleSplit：此模塊創建單個訓練/測試組具有同樣的平衡（分層）的類。基本上這就是你想要的n_iter=1。你可以提測試尺寸這裏一樣train_test_split

代碼：

>>> sss = StratifiedShuffleSplit(y, n_iter=1, test_size=0.5, random_state=0) 
>>> len(sss) 
1 
>>> for train_index, test_index in sss: 
... print("TRAIN:", train_index, "TEST:", test_index) 
... X_train, X_test = X[train_index], X[test_index] 
... y_train, y_test = y[train_index], y[test_index] 
>>> # fit and predict with your classifier using the above X/y train/test 

Answer 3

下面是連續/迴歸數據爲例（直到this issue on GitHub解決）。

# Your bins need to be appropriate for your output values 
# e.g. 0 to 50 with 25 bins 
bins  = np.linspace(0, 50, 25) 
y_binned = np.digitize(y_full, bins) 
X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y_binned) 

Answer 4

除了由@Andreas穆勒接受的答案，只是想補充一點，作爲@tangy上面提到的：

StratifiedShuffleSplit最接近train_test_split（分層= Y） 用的附加功能：

1. 分層默認
2. 通過指定n_splits，它一再會將數據

Answer 5

#train_size is 1 - tst_size - vld_size 
tst_size=0.15 
vld_size=0.15 

X_train_test, X_valid, y_train_test, y_valid = train_test_split(df.drop(y, axis=1), df.y, test_size = vld_size, random_state=13903) 

X_train_test_V=pd.DataFrame(X_train_test) 
X_valid=pd.DataFrame(X_valid) 

X_train, X_test, y_train, y_test = train_test_split(X_train_test, y_train_test, test_size=tst_size, random_state=13903)