Stratified Train / Test-split ใน scikit-learn

ฉันต้องการแบ่งข้อมูลของฉันออกเป็นชุดฝึก (75%) และชุดทดสอบ (25%) ขณะนี้ฉันใช้รหัสด้านล่าง:

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

อย่างไรก็ตามฉันต้องการแบ่งชั้นชุดข้อมูลการฝึกอบรมของฉัน ฉันจะทำอย่างไร? ฉันได้ตรวจสอบStratifiedKFoldวิธีการนี้แล้ว แต่อย่าให้ฉันระบุการแบ่ง 75% / 25% และแบ่งชั้นชุดข้อมูลการฝึกอบรมเท่านั้น

[อัปเดตสำหรับ 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]

มีคำขอดึงเป็นที่นี่ แต่คุณสามารถทำtrain, test = next(iter(StratifiedKFold(...))) และใช้รถไฟและทดสอบดัชนีได้หากต้องการ

TL; DR: ใช้StratifiedShuffleSplitกับtest_size=0.25

Scikit-learn มีสองโมดูลสำหรับ Stratified Splitting:

1. StratifiedKFold : โมดูลนี้มีประโยชน์ในฐานะตัวดำเนินการตรวจสอบความถูกต้องข้ามแบบ k-fold โดยตรงเนื่องจากจะตั้งค่าชุด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

2. 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

คุณสามารถทำได้ด้วยtrain_test_split()วิธีการที่มีอยู่ใน Scikit learn:

from sklearn.model_selection import train_test_split 
train, test = train_test_split(X, test_size=0.25, stratify=X['YOUR_COLUMN_LABEL']) 

ฉันได้เตรียม GitHub Gist สั้น ๆ ซึ่งแสดงให้เห็นว่าstratifyตัวเลือกทำงานอย่างไร:

https://gist.github.com/SHi-ON/63839f3a3647051a180cb03af0f7d0d9

นี่คือตัวอย่างสำหรับข้อมูลต่อเนื่อง / ถดถอย (จนกว่าปัญหานี้ใน GitHubจะได้รับการแก้ไข)

min = np.amin(y)
max = np.amax(y)

# 5 bins may be too few for larger datasets.
bins     = np.linspace(start=min, stop=max, num=5)
y_binned = np.digitize(y, bins, right=True)

X_train, X_test, y_train, y_test = train_test_split(
    X, 
    y, 
    stratify=y_binned
)

• ที่ไหนstartเป็นนาทีstopค่าสูงสุดของเป้าหมายต่อเนื่องของคุณ
• หากคุณไม่ได้ตั้งค่าright=Trueจะทำให้ค่าสูงสุดของคุณเป็นถังขยะแยกกันมากขึ้นหรือน้อยลงและการแยกของคุณจะล้มเหลวเสมอเพราะจะมีตัวอย่างน้อยเกินไปในถังพิเศษนั้น

นอกจากคำตอบที่ได้รับการยอมรับจาก @Andreas Mueller แล้วเพียงแค่ต้องการเพิ่มเป็น @tangy ที่กล่าวถึงข้างต้น:

StratifiedShuffleSplitใกล้เคียงที่สุดกับtrain_test_split (stratify = y) พร้อมคุณสมบัติเพิ่มเติมของ:

1. แบ่งชั้นตามค่าเริ่มต้น
2. โดยการระบุn_splitsมันจะแยกข้อมูลซ้ำ ๆ

#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)

การอัปเดตคำตอบ @tangy จากด้านบนเป็นเวอร์ชันปัจจุบันของ scikit-learn: 0.23.2 ( เอกสาร StratifiedShuffleSplit )

from sklearn.model_selection import StratifiedShuffleSplit

n_splits = 1  # We only want a single split in this case
sss = StratifiedShuffleSplit(n_splits=n_splits, test_size=0.25, random_state=0)

for train_index, test_index in sss.split(X, y):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]