1. 程式人生 > >sklearn資料集分割函式[StratifiedShuffleSplit、split]

sklearn資料集分割函式[StratifiedShuffleSplit、split]

  • sklearn.model_selection.StratifiedShuffleSplit(n_splits=10, test_size=’default’, train_size=None, random_state=None)
    將資料集劃分為幾組資料集/測試集。
  • split(X, y)
    按照y的值將資料集分為訓練集或測試集,保證訓練集和測試集中各類y值所佔的比例與原資料集相同。

考慮一個場景,房屋資料集中認為地區收入在房價中佔很重要的作用,想分開的訓練集和測試集中各種收入所佔的比重與原資料集中所佔的比重相同。
程式碼實現:

    from sklearn.model_selection import
StratifiedShuffleSplit split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42) for train_index, test_index in split.split(housing, housing["income_cat"]): strat_train_set = housing.loc[train_index] strat_test_set = housing.loc[test_index]