自定义Sklearn CV拆分器遇索引越界错误，寻求技术帮助

问题背景 #

构建自定义CV拆分器GroupShuffleTwoColumnsSplit时遭遇索引越界错误，当前split函数仅实现简单随机拆分逻辑用于排查，但错误仍存在。核心目标：

1. 先将全量数据拆分为开发集（Dev）和测试集（Test）
2. 将拆分器传入RandomizedSearchCV，从开发集中拆分训练集和验证集

已执行操作 #

• 定义包含n_splits、train_size、test_size、random_state属性的自定义拆分器类
• 编写split迭代器方法，随机选择索引返回训练/测试索引
• 生成测试数据
• 首次初始化拆分器，通过next()获取Dev/Test集索引
• 再次初始化拆分器，传入RandomizedSearchCV用于Train/Val集拆分
• 基于线性回归定义随机搜索超参数空间
• 将自定义拆分器传入RandomizedSearchCV

尝试过的排查方案 #

• 替换为内置GroupShuffleSplit：运行正常
• 跳过Dev-Test拆分，直接将全量数据传入RandomizedSearchCV：运行正常
• 仅初始化一次拆分器同时用于两步拆分：仍报索引越界错误

问题代码 #

import pandas as pd
import numpy as np
from sklearn.model_selection import GroupShuffleSplit
import random
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import RandomizedSearchCV


class GroupShuffleTwoColumnsSplit:

    def __init__(self, n_splits, train_size=None, test_size=None, random_state=None):
        self.n_splits = n_splits
        self.train_size = train_size
        self.test_size = test_size
        self.random_state = random_state

        if self.train_size is None:
            self.train_size = 1 - self.test_size
        if self.test_size:
            self.train_size = 1 - self.test_size

    def split(self, X, y, groups=None):

        series_0 = groups.iloc[:, 0]

        for n in range(self.n_splits):

            indices = series_0.index.tolist()

            ratio = 1 - self.train_size
            num_elements = int(len(indices) * ratio)
            indices_test = random.sample(indices, num_elements)
            indices_train = [x for x in indices if x not in indices_test]
            yield indices_train, indices_test

    def get_n_splits(self, X, y, groups=None):
        return self.n_splits


multiplicator = 10
list_values_1 = [831, 832, 833, 834, 835]
list_values_2 = [1, 2, 3, 4, 5]
col_1 = np.repeat(list_values_1, multiplicator)
col_2 = np.tile(list_values_2, multiplicator)
y_col = np.random.rand(len(list_values_1) * multiplicator)

df = pd.DataFrame({"serie": col_1, "poles": col_2, "y": y_col})
X = df.loc[:, ["serie", "poles"]]
y = df.loc[:, "y"]

### 初始Dev-Test拆分
group_columns = df.loc[:, ["serie", "poles"]]
splitter = GroupShuffleTwoColumnsSplit(n_splits=1, train_size=0.8)

dev_id, test_id = next(splitter.split(X, y, group_columns))
X_dev = df.loc[dev_id, :]
X_test = df.loc[test_id, :]
y_dev = df.loc[dev_id]
y_test = df.loc[test_id]

### 超参数搜索
group_columns = X_dev.loc[:, ["serie", "poles"]]
splitter = GroupShuffleTwoColumnsSplit(n_splits=5, train_size=0.8)

param_dist = {
    "fit_intercept": [True, False],
}

random_search = RandomizedSearchCV(
    estimator=LinearRegression(),
    param_distributions=param_dist,
    n_iter=10,
    cv=splitter,
)

random_search.fit(X_dev, y_dev, groups=group_columns)

错误原因分析 #

1. 索引类型不匹配：Sklearn的CV拆分器要求返回位置索引（0-based整数），但自定义拆分器返回的是原DataFrame的标签索引。当传入X_dev（原df的子集，保留原标签索引）时，拆分器生成的索引超出了X_dev的位置范围，导致索引越界。
2. 初始化逻辑漏洞：__init__中若同时传入train_size和test_size，train_size会被重复覆盖；且未处理两者都未传入的边界情况，可能导致后续比例计算错误。
3. 随机种子未固定：使用random.sample但未绑定随机种子，拆分结果不可复现，不利于调试。

修复后的代码 #

import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import RandomizedSearchCV


class GroupShuffleTwoColumnsSplit:

    def __init__(self, n_splits, train_size=None, test_size=None, random_state=None):
        self.n_splits = n_splits
        self.train_size = train_size
        self.test_size = test_size
        self.random_state = random_state
        self.rng = np.random.RandomState(random_state)

        # 修复初始化逻辑：优先处理train_size，再处理test_size，避免重复覆盖
        if self.train_size is None and self.test_size is None:
            raise ValueError("必须指定train_size或test_size中的一个")
        if self.test_size is not None:
            self.train_size = 1 - self.test_size
        elif self.train_size is None:
            self.train_size = 1 - self.test_size

    def split(self, X, y, groups=None):
        # 使用位置索引而非标签索引，适配Sklearn要求
        n_samples = len(groups)
        indices = np.arange(n_samples)

        for _ in range(self.n_splits):
            self.rng.shuffle(indices)
            split_idx = int(self.train_size * n_samples)
            indices_train = indices[:split_idx]
            indices_test = indices[split_idx:]
            yield indices_train, indices_test

    def get_n_splits(self, X, y, groups=None):
        return self.n_splits


# 生成测试数据
multiplicator = 10
list_values_1 = [831, 832, 833, 834, 835]
list_values_2 = [1, 2, 3, 4, 5]
col_1 = np.repeat(list_values_1, multiplicator)
col_2 = np.tile(list_values_2, multiplicator)
y_col = np.random.rand(len(list_values_1) * multiplicator)

df = pd.DataFrame({"serie": col_1, "poles": col_2, "y": y_col})
X = df.loc[:, ["serie", "poles"]]
y = df.loc[:, "y"]

# 初始Dev-Test拆分
group_columns = df.loc[:, ["serie", "poles"]]
splitter = GroupShuffleTwoColumnsSplit(n_splits=1, train_size=0.8, random_state=42)
dev_id, test_id = next(splitter.split(X, y, group_columns))
X_dev = df.iloc[dev_id, :]  # 改用iloc按位置索引取值
X_test = df.iloc[test_id, :]
y_dev = df.iloc[dev_id, 2]
y_test = df.iloc[test_id, 2]

# 超参数搜索
group_columns_dev = X_dev.loc[:, ["serie", "poles"]]
splitter_cv = GroupShuffleTwoColumnsSplit(n_splits=5, train_size=0.8, random_state=42)

param_dist = {
    "fit_intercept": [True, False],
}

random_search = RandomizedSearchCV(
    estimator=LinearRegression(),
    param_distributions=param_dist,
    n_iter=10,
    cv=splitter_cv,
    random_state=42
)

random_search.fit(X_dev, y_dev, groups=group_columns_dev)
print(random_search.best_params_)

关键修改点 #

• 索引类型修正：split函数返回位置索引（np.arange(n_samples)），Dev-Test拆分时改用iloc按位置取值，避免标签索引冲突
• 初始化逻辑修复：明确train_size和test_size的优先级，添加边界值校验
• 随机种子固定：使用np.random.RandomState绑定随机种子，保证拆分结果可复现

内容的提问来源于stack exchange，提问作者user23501795