分層拆分

分層拆分保證拆分後的數據集標籤列比例還一樣。比如在原來數據集中正負樣本比例是2:1,那麼在拆分後的測試集和訓練集中,正負標籤也是2:1。
可以用來修正隨機拆分後的測試集和訓練中比例不一樣的問題。
如果正樣本特別少,並且測試集也很少,那麼測試集有可能抽不到正樣本,可以使用分層採樣。

使用sklearn 測試,不使用分層:

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.DataFrame(
    data={
        'c1': ['a', 'b', 'c', 'd', 'e', 'f'],
        'label': [1, 1, 1, 1, 0, 0]
    }
)
# X_train, X_test = train_test_split(df, test_size=0.333, random_state=100, stratify=df['label'])  # 使用分層抽樣,指定分層抽樣依據的列
X_train, X_test = train_test_split(df, test_size=0.333, random_state=100)

print('X_train: ')
print(X_train)
print('X_test: ')
print(X_test)
X_train: 
  c1  label
4  e      0
3  d      1
5  f      0
0  a      1
X_test: 
  c1  label
1  b      1
2  c      1

使用分層抽樣的輸出:

X_train: 
  c1  label
1  b      1
3  d      1
5  f      0
2  c      1
X_test: 
  c1  label
4  e      0
0  a      1
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章