sklearn进行StratifiedShuffleSplit时的index问题

1. 问题的提出

sklearn中常采用StratifiedShuffleSplit对样本进行分层采样，其返回数据集划分后对应训练集和测试集的index。如果我们的样本数据类型为pandas.DataFrame，其也自带index属性，那这两个index有无联系呢？

2. 问题的分析

下文为一个简单的测试：

from sklearn.model_selection import StratifiedShuffleSplit
import pandas as pd
import numpy as np

# 1. 构造数据
index = np.arange(100)    # DataFrame的index
np.random.seed(0)
np.random.shuffle(index)     # 随机打乱
cla = np.random.choice([0, 1], p=[0.2, 0.8], size=100)    # 类别数据

# 2. 划分数据
sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=64)

for train_index, test_index in sss.split(df, df['b']):

    train_data = df.loc[train_index]
    test_data = df.loc[test_index]
    print(train_data['b'].sum())    # 返回68，符合预期

返回样本中的正、负例数量看似正常且符合预期。如果我们把原始数据的index修改下呢？

# 1. 构造数据
index = np.arange(100, 200)    # DataFrame的index
np.random.seed(0)
np.random.shuffle(index)     # 随机打乱
cla = np.random.choice([0, 1], p=[0.2, 0.8], size=100)    # 类别数据

# 2. 划分数据
sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=64)

for train_index, test_index in sss.split(df, df['b']):

    train_data = df.loc[train_index]
    test_data = df.loc[test_index]
    print(train_data['b'].sum())    

程序会直接报错：

KeyError: "None of [Int64Index([73, 64, 52, 65, 38, 63, 67, 84, 24, 14, 83, 48, 44, 25, 39, 58, 91,\n            85, 22, 70, 28, 78, 89,  4, 66, 97, 69, 27,  0, 19,  7, 34, 95,  8,\n            33,  3, 31, 16, 32,  6, 87, 68, 98, 43, 57, 72, 82, 59, 37, 56, 76,\n            61, 49, 12, 88, 50, 81, 30, 36, 60, 15,  1, 92, 40, 99, 77, 62, 10,\n            80, 51, 45, 11, 94, 41, 96, 55,  5, 18,  2, 53],\n           dtype='int64')] are in the [index]"

这提醒我们，StratifiedShuffleSplit后的index并不能在原始的pandas.DataFrame数据中找到！

回过头看下类StratifiedShuffleSplit在split时的输入参数说明：

# split方法的参数
X : array-like, shape (n_samples, n_features)

这说明，StratifiedShuffleSplit在实际处理时，会将其转换类似于numpy.ndarray这样的数据，其不包含index这样的属性的。所以split的实质只是对array行号索引的记录，与pandas.DataFrame无任何关系。

上面的案例第一次运行看似正常，其实只是因为pandas.DataFrame的index和转换后array的index取值范围恰巧是一致的！

这时有人会想到，既然这样那我们每次在划分前re_index，将其强制转换为与array长度相当的索引不就可以了吗？

实际上，这种方法是有问题和缺陷的！
一方面，array的index是选取的绝对行号，而若使用loc用的是’index’名称，两者并不一样，可能达不到分层采样的目的。因此，此时必须配合iloc使用。

另一方面，index往往是作为数据样本的标记，而reset_index会打乱这种标记，使得后期在样本查找时更加麻烦。

3. 另一个隐晦的易错点

如上节所述，将loc改为iloc这种采取绝对位置的切片方式进行数据划分，我们引入一个更加复杂的划分“训练集-验证集-测试集”的示例：

from sklearn.model_selection import StratifiedShuffleSplit
import pandas as pd
import numpy as np

# 1. 构造数据
index = np.arange(100)    # DataFrame的index
np.random.seed(0)
np.random.shuffle(index)     # 随机打乱
cla = np.random.choice([0, 1], p=[0.2, 0.8], size=100)    # 类别数据

# 2.划分训练集和测试集
for train_index, test_index in sss.split(df, df['b']):
    train_data = df.iloc[train_index]
    test_data = df.iloc[test_index]

    # 将训练集进一步划分为训练数据和验证数据
    for train_index_, validation_index_ in sss.split(train_data['b'], train_data['b']):
        train_data_ = df.iloc[train_index_]
        validate_data_ = df.iloc[validation_index_]

        train_index_set = set(train_data_.index)
        validate_index_set = set(validate_data_.index)
        test_index_set = set(test_data.index)

        if train_index_set.intersection(validate_index_set) or train_index_set.intersection(test_index_set) or validate_index_set.intersection(test_index_set):
            print("发生数据交叉")
        else:
            print("数据被隔离划分")

从输出结果来看，竟然发生了数据交叉，也就是说训练集、验证集和测试集之间出现了共同数据，这是数据分析所必须避免的问题。

其原因在于训练集和验证集划分时，对全量数据（包括测试数据）进行了iloc，会导致部分数据同时存在与测试集和训练集/验证集中。因此，训练集和验证集划分时必须基于先前划分后的训练数据基础上。修改后的代码为：

# 划分训练集和测试集
for train_index, test_index in sss.split(df, df['b']):
    train_data = df.iloc[train_index]
    test_data = df.iloc[test_index]

    # 将训练集进一步划分为训练数据和验证数据
    for train_index_, validation_index_ in sss.split(train_data['b'], train_data['b']):
        train_data_ = train_data.iloc[train_index_]    # 此时必须为train_data
        validate_data_ = train_data.iloc[validation_index_]   # 此时必须为train_data

        train_index_set = set(train_data_.index)
        validate_index_set = set(validate_data_.index)
        test_index_set = set(test_data.index)

        if train_index_set.intersection(validate_index_set) or train_index_set.intersection(test_index_set) or validate_index_set.intersection(test_index_set):
            print("发生数据交叉")
        else:
            print("数据被隔离划分")

4. 总结

本文来源于对平时使用sklearn时遇到的一次偶然错误。在进行数据集划分建议对数据进行划分结果进行校验，以保证训练集、验证集和测试集的绝对隔离，以及分层采样的合理性。

建议的做法为：
方案一: 若运用StratifiedShuffleSplit进行划分，输入数据为DataFrame时，必须注意两者index的差别，采用iloc的切片方式。
方案二: 在这种场景下，更适合用train_test_split的直接得到划分后的数据，其默认采用分层采样方法。