pip安装iterstrat.ml_stratifiers import MultilabelStratifiedKFold, MultilabelStratifiedShuffleSplit

安装方法

pip install iterative-stratification

开源源码

介绍

MultilabelStratifiedKFold是用于多标签多分类的K折交叉验证

from iterstrat.ml_stratifiers import MultilabelStratifiedKFold, MultilabelStratifiedShuffleSplit
# 5折
nfold = 5
kf = MultilabelStratifiedKFold(n_splits=nfold, shuffle=True, random_state=2020)

lr_oof = np.zeros(label.shape)
# 预测结果每个label的概率
lr_predictions = np.zeros((len(test), label.shape[1]))

i = 0
for train_index, valid_index in kf.split(train_df, label):
    X_train, label_train = train_df[train_index], label[train_index]
    X_valid, label_valid = train_df[valid_index], label[valid_index]

    base = LogisticRegression(C=1)
    model = OneVsRestClassifier(base, n_jobs=20)
    model.fit(X_train, label_train)
    # 返回label的概率
    probality = model.predict_proba(X_valid,)
    print(probality)