sklearn.cross_validation.StratifiedShuffleSplit - 错误:"索引超出范围"

时间:2015-05-04 06:35:04

标签: python pandas scikit-learn

我试图使用Scikit-learn的Stratified Shuffle Split来分割样本数据集。我按照Scikit-learn文档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

有人可以指出我在这里做错了什么吗?谢谢!

1 个答案:

答案 0 :(得分:45)

您正在遇到Pandas DataFrame索引与NumPy ndarray索引的不同约定。数组train_indextest_index是行索引的集合。但是data是一个Pandas DataFrame对象,当您在该对象中使用单个索引时,就像在data[train_index]中一样,Pandas希望train_index包含列< / em>标签而不是行索引。您可以使用.values

将数据帧转换为NumPy数组
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]

或使用Pandas .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类型的xtestDataFrame而不是ndarray,因此保留了列标签。