使用Sklearn的ShuffleSplit拆分数据集时遇KeyError 223问题求助

我来帮你搞定这个头疼的KeyError问题～其实原因很明确：ShuffleSplit的split()方法返回的是数据的位置索引（也就是从0开始的连续整数，对应每一行在数据集里的位置），但如果你的df_noyau_yes的索引不是连续的0起始整数（比如之前筛选过数据导致索引断档、或者用了自定义的字符串/非连续数字索引），直接用这些返回的索引去取数（比如用df.loc[train_index]）就会找不到对应的键，自然就抛出KeyError了。

给你两种靠谱的解决方案：

方法一：用iloc按位置提取数据 #

iloc是Pandas里专门按位置取数的方法，正好匹配ShuffleSplit返回的位置索引，直接用就不会有问题：

from sklearn.model_selection import ShuffleSplit

rs = ShuffleSplit(n_splits=2, test_size=.25, random_state=0)
for train_index, test_index in rs.split(df_noyau_yes):
    # 用iloc提取训练集和测试集
    X_train = df_noyau_yes.iloc[train_index]
    X_test = df_noyau_yes.iloc[test_index]
    # 喂入自编码器训练（自编码器通常输入输出一致）
    # your_autoencoder.fit(X_train, X_train)

方法二：重置DataFrame的索引 #

如果习惯用loc取数，可以先把数据集的索引重置为连续的0起始整数，之后就可以正常用返回的索引取数了：

from sklearn.model_selection import ShuffleSplit

# 重置索引，drop=True会丢弃原来的旧索引列
df_reset = df_noyau_yes.reset_index(drop=True)
rs = ShuffleSplit(n_splits=2, test_size=.25, random_state=0)
for train_index, test_index in rs.split(df_reset):
    X_train = df_reset.loc[train_index]
    X_test = df_reset.loc[test_index]
    # 进行自编码器训练

额外小提示 #

如果你只需要一次拆分（比如n_splits设为1），其实可以直接用train_test_split，它会直接返回拆分好的数据集，不用手动处理索引，更省心：

from sklearn.model_selection import train_test_split

X_train, X_test = train_test_split(df_noyau_yes, test_size=0.25, random_state=0)
# 直接用拆分好的数据训练自编码器即可

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