Phân tầng Train / Test-split trong scikit-learning

Tôi cần chia dữ liệu của mình thành tập huấn luyện (75%) và tập kiểm tra (25%). Tôi hiện đang làm điều đó với mã bên dưới:

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

Tuy nhiên, tôi muốn phân tầng tập dữ liệu đào tạo của mình. Làm thế nào để làm điều đó? Tôi đã xem xét StratifiedKFoldphương pháp, nhưng không cho phép tôi xác định tỷ lệ 75% / 25% và chỉ phân tầng tập dữ liệu đào tạo.

[cập nhật cho 0.17]

Xem tài liệu của 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)

[/ cập nhật cho 0.17]

Có một yêu cầu kéo ở đây . Nhưng bạn chỉ có thể làm train, test = next(iter(StratifiedKFold(...))) và sử dụng các chỉ số đào tạo và kiểm tra nếu bạn muốn.

TL; DR: Sử dụng StratifiedShuffleSplit vớitest_size=0.25

Scikit-learning cung cấp hai mô-đun cho Phân tách phân tầng:

1. StratifiedKFold : Mô-đun này hữu ích như một toán tử xác thực chéo k-lần trực tiếp: vì nó sẽ thiết lập các n_foldstập huấn luyện / kiểm tra sao cho các lớp đều cân bằng trong cả hai.

Đây là một số mã (trực tiếp từ tài liệu trên)

>>> 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 : Mô-đun này tạo một tập huấn luyện / kiểm tra duy nhất có các lớp (phân tầng) cân bằng như nhau. Về cơ bản đây là những gì bạn muốn với n_iter=1. Bạn có thể đề cập đến kích thước thử nghiệm ở đây giống như trongtrain_test_split

Mã:

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

Đây là ví dụ về dữ liệu hồi quy / liên tục (cho đến khi vấn đề này trên GitHub được giải quyết).

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
)

• Đâu startlà tối thiểu và stoplà tối đa của mục tiêu liên tục của bạn.
• Nếu bạn không đặt right=Truethì nó sẽ ít nhiều khiến giá trị tối đa của bạn trở thành một thùng riêng biệt và việc phân chia của bạn sẽ luôn không thành công vì có quá ít mẫu trong thùng phụ đó.

Bạn có thể làm điều đó một cách đơn giản với train_test_split()phương pháp có sẵn trong Scikit learning:

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

Tôi cũng đã chuẩn bị một GitHub Gist ngắn cho thấy cách stratifyhoạt động của tùy chọn:

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

Ngoài câu trả lời được chấp nhận bởi @Andreas Mueller, chỉ muốn thêm câu trả lời đó như @tangy đã đề cập ở trên:

StratifiedShuffleSplit gần giống nhất với train_test_split ( stratify = y) với các tính năng bổ sung của:

1. phân tầng theo mặc định
2. bằng cách chỉ định n_splits , nó liên tục chia nhỏ dữ liệu

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

Đang cập nhật câu trả lời @tangy từ bên trên lên phiên bản hiện tại của scikit-learning: 0.23.2 ( tài liệu 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]