如何用sklearn实现基于分层抽样的交叉验证（构造一个类）

import numpy as np
from sklearn.metrics import f1_score
from sklearn.model_selection import StratifiedKFold
from sklearn.base import clone, BaseEstimator, TransformerMixin

class stratified_cross_val_score(BaseEstimator, TransformerMixin):
    """实现基于分层抽样的k折交叉验证"""
    
    def __init__(self, model, random_state=0, cv=5, pattern="classification"):
        """
        :model: 训练的模型（回归或分类）
        :random_state: 模型的随机种子值
        :cv: 交叉验证的次数
        :pattern: classification和regression两种选择
        """
        self.model = model
        self.random_state = random_state
        self.cv = cv
        self.pattern = pattern
        self.scores_ = []  
        self.best_score_ = []
        self.estimators_ = []
        self.best_estimator_ = []
        
        self.i = 0            
    
    def fit(self, X, y, layer_tag):
        """
        :param X: 只含有特征值的完整数据集
        :param y: 只含有标签值的完整数据集
        :param tag: 只含有分层依据的完整数据集（此例是KMeans聚类结果）
        
        """
        skfolds = StratifiedKFold(n_splits=self.cv, random_state=self.random_state, shuffle=True)

        for train_index, test_index in skfolds.split(X, layer_tag):
            # 复制要训练的模型（分类或回归）
            clone_model = clone(self.model)
            strat_X_train_folds, strat_X_test_fold = X.iloc[train_index], X.iloc[test_index]
            strat_y_train_folds, strat_y_test_fold = y.iloc[train_index], y.iloc[test_index]
            
            # 训练模型
            clone_model.fit(strat_X_train_folds, strat_y_train_folds)
            # 保留模型
            self.estimators_.append(clone_model)
            # 预测值（这里是分类模型的分类结果）
            test_labels_pred = clone_model.predict(strat_X_test_fold)
            
            if self.pattern == "classification":
                # 分类模型用F1值
                score_fold = f1_score(y.iloc[test_index], test_labels_pred, average="weighted")
            elif self.pattern == "regression":
                # 回归模型使用r2值
                score_fold = r2_score(y.iloc[test_index], test_labels_pred)
            
            # 避免重复向列表里重复添加值
            if self.i < self.cv:
                self.scores_.append(score_fold)
            else:
                None
                
            self.i += 1
        
        # 获取评分最高模型的索引
        argmax  = np.argmax(self.scores_)
        self.best_score_ = self.scores_[argmax]
        self.best_estimator_ = self.estimators_[argmax]
        
    def transform(self, X, y=None):
        return self
    
    def mean(self):
        """返回交叉验证评分的平均值"""
        return np.array(self.scores_).mean()
    
    def std(self):
        """返回交叉验证评分的标准差"""
        return np.array(self.scores_).std()

from sklearn.linear_model import SGDClassifier

# 分类模型
clf_model = SGDClassifier(max_iter=5, tol=-np.infty, random_state=42)
# 基于分层抽样的交叉验证，data是只含特征值的完整数据集，labels是只含标签值的完整数据集
clf_cross_val = stratified_cross_val_score(clf_model, data, labels, cv=5, random_state=42)
# data2是含有特征值和聚类结果的完整数据集
clf_cross_val.fit(data2, data2["km_clustering_label"])

# 每折交叉验证的评分
clf_cross_val.score

[0.26987105798643873,
 0.26073985257088084,
 0.26210674310514226,
 0.2737537146865997,
 -0.06799899887753336]

# 五折交叉验证中最好的评分
clf_cross_val.best_score_

0.2737537146865997

# 交叉验证评分的平均值
clf_cross_val.mean()

0.1996944738943056

# 交叉验证评分的标准差
clf_cross_val.std()

0.13393367692408878