scikit-learn - 是否可以在分层KFold 中取回列表？

Question

我想做这样的事情：

Skf = sklearn.model_selection.StratifiedKFold(n_splits = 5, shuffle = True) 
ALPHA,BETA  = Skf.split(data_X, data_Y)

接着 ：

for train_index, test_index in ALPHA,BETA

但是，它不起作用，为什么以及如何绕过该问题？

我的想法是我想在代码的不同部分多次使用相同的拆分...我不知道如何“存储”拆分。

Answer 1

是的你可以。您可以指定随机数生成器使用的种子，以便在不同的运行中获得相同的拆分。只需指定random_state参数！

SEED = 42
Skf = sklearn.model_selection.StratifiedKFold(n_splits=5,
                                              shuffle=True,
                                              random_state=SEED)