sklearn.cross_validation.StratifiedShuffleSplit - 错误：“索引超出范围”

Question

我试图使用Scikit-learn的分层随机拆分拆分示例数据集。我也跟着上所示的例子Scikit学习文档here

import pandas as pd 
import numpy as np 
# UCI's wine dataset 
wine = pd.read_csv("https://s3.amazonaws.com/demo-datasets/wine.csv") 

# separate target variable from dataset 
target = wine['quality'] 
data = wine.drop('quality',axis = 1) 

# Stratified Split of train and test data 
from sklearn.cross_validation import StratifiedShuffleSplit 
sss = StratifiedShuffleSplit(target, n_iter=3, test_size=0.2) 

for train_index, test_index in sss: 
    xtrain, xtest = data[train_index], data[test_index] 
    ytrain, ytest = target[train_index], target[test_index] 

# Check target series for distribution of classes 
ytrain.value_counts() 
ytest.value_counts() 

然而，一旦运行该脚本，我得到以下错误：

IndexError: indices are out-of-bounds 

可能有人请指出我在做什么错在这里？谢谢！

Answer 1

您正在运行Pandas DataFrame索引与NumPy ndarray索引的不同约定。阵列train_index和是行索引的集合。但是data是Pandas DataFrame对象，并且当您使用单个索引导入该对象时，如在data[train_index]中，Pandas期望train_index包含列列标签而非行索引。您可以将数据帧转换为NumPy的阵列，使用.values：

data_array = data.values 
for train_index, test_index in sss: 
    xtrain, xtest = data_array[train_index], data_array[test_index] 
    ytrain, ytest = target[train_index], target[test_index] 

或使用熊猫.iloc访问：

for train_index, test_index in sss: 
    xtrain, xtest = data.iloc[train_index], data.iloc[test_index] 
    ytrain, ytest = target[train_index], target[test_index] 

我倾向于第二种方法，因为它提供了xtrain和xtest键入DataFrame而不是ndarray，并保留列标签。