Scikit-learn中的train_test_split和indices

90

当使用train_test_split()时,我如何获取数据的原始索引?

我手头有以下内容:

from sklearn.cross_validation import train_test_split
import numpy as np
data = np.reshape(np.randn(20),(10,2)) # 10 training examples
labels = np.random.randint(2, size=10) # 10 labels
x1, x2, y1, y2 = train_test_split(data, labels, size=0.2)

但这并不提供原始数据的索引。一种解决方法是将索引添加到数据中(例如data = [(i, d) for i, d in enumerate(data)]),然后将它们传递给train_test_split,再次展开。还有更简洁的解决方案吗?


3
请注意 sklearn.model_selection.ShuffleSplitsklearn.model_selection.StratifiedShuffleSplit。它们是机器学习中常用的交叉验证方法,可用于评估模型性能并选择最佳参数。其中ShuffleSplit是将数据集随机划分为训练集和测试集,而StratifiedShuffleSplit则通过保持类别比例来进行随机划分。 - Jost
5个回答

132

正如Julien所说,您可以使用pandas数据帧或系列,但如果您想限制自己只使用numpy,则可以传递一个额外的索引数组:

from sklearn.model_selection import train_test_split
import numpy as np
n_samples, n_features, n_classes = 10, 2, 2
data = np.random.randn(n_samples, n_features)  # 10 training examples
labels = np.random.randint(n_classes, size=n_samples)  # 10 labels
indices = np.arange(n_samples)
(
    data_train,
    data_test,
    labels_train,
    labels_test,
    indices_train,
    indices_test,
) = train_test_split(data, labels, indices, test_size=0.2)

1
截至NumPy v1.1,第三行应为data = np.reshape(np.random.randn(20),(10,2)); 最后一行应为... train_test_split(data, labels, indices, test_size=0.2) - pepe
17
这应该是被接受的回答,因为它并没有使用任何额外的包,而是使用了sklearn。它比pandas更加控制所发生的事情。 - Irene
@ogrisel 你好,我有一个类似的问题,能否请您检查一下 https://stackoverflow.com/questions/48734942/how-to-maintain-the-x-axis-value-from-the-train-test-split-when-plotting-y-train 谢谢。 - user9238790
嘿,在你的第三行中,n_class是什么意思?“class”的数量,你指的是什么?提前感谢。 - Yun Tae Hwang
这是要生成随机分类标签的目标标签类数,存储在 labels 变量中。 - ogrisel
如果添加的索引不止一个,该代码也可以正常工作。例如,对于需要拆分的token_lists:x1,x2,y1,y2,idx1,idx2,t1,t2 = train_test_split(data,labels,indices,token_lists,test_size = 0.2) - questionto42

59

Scikit learn和Pandas的结合非常出色,因此我建议您使用它。这里有一个例子:

In [1]: 
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
data = np.reshape(np.random.randn(20),(10,2)) # 10 training examples
labels = np.random.randint(2, size=10) # 10 labels

In [2]: # Giving columns in X a name
X = pd.DataFrame(data, columns=['Column_1', 'Column_2'])
y = pd.Series(labels)

In [3]:
X_train, X_test, y_train, y_test = train_test_split(X, y, 
                                                    test_size=0.2, 
                                                    random_state=0)

In [4]: X_test
Out[4]:

     Column_1    Column_2
2   -1.39       -1.86
8    0.48       -0.81
4   -0.10       -1.83

In [5]: y_test
Out[5]:

2    1
8    1
4    1
dtype: int32
您可以直接在DataFrame/Series上调用任何scikit函数,它们都可以正常工作。
举个例子,假设您想要执行LogisticRegression,以下是如何以优美的方式检索系数:
In [6]: 
from sklearn.linear_model import LogisticRegression

model = LogisticRegression()
model = model.fit(X_train, y_train)

# Retrieve coefficients: index is the feature name (['Column_1', 'Column_2'] here)
df_coefs = pd.DataFrame(model.coef_[0], index=X.columns, columns = ['Coefficient'])
df_coefs
Out[6]:
            Coefficient
Column_1    0.076987
Column_2    -0.352463

1
此外,看起来您的问题可能是代码存在问题或者您正在使用已弃用的scikit和numpy版本(np.randn在我的版本中不存在,并且train_test_split中使用的是test_size而不是size)。 - Julien Marrec
2
我编辑了我的答案,展示了如何使用来自pandas数据框的特征名称检索系数。这可能会在未来为您节省一些时间。 - Julien Marrec
嗨 @Julien Marrec,我尝试了您的解决方案,但它没有起作用,您可以在这里检查一下吗 https://stackoverflow.com/questions/48734942/how-to-maintain-the-x-axis-value-from-the-train-test-split-when-plotting-y-train - user9238790
4
如果我已经在创建索引之前分割了我的数据,怎么办? - Samuel Nde
1
这个回答如何回答原始问题?看起来你只是创建了一个新的数据框,并将原始数据框的索引应用于新创建的数组。当你使用训练测试拆分随机化数据时,你正在洗牌行,并且仅仅将先前数据框的索引应用于洗牌后的数据并不能准确地访问原始数据的索引。我有什么遗漏吗? - devdrc
显示剩余3条评论

12

这是最简单的解决方案(在另一个答案中,Jibwa让它看起来很复杂),而不必自己生成索引 - 只需使用ShuffleSplit对象生成1个拆分即可。

import numpy as np 
from sklearn.model_selection import ShuffleSplit # or StratifiedShuffleSplit
sss = ShuffleSplit(n_splits=1, test_size=0.1)

data_size = 100
X = np.reshape(np.random.rand(data_size*2),(data_size,2))
y = np.random.randint(2, size=data_size)

sss.get_n_splits(X, y)
train_index, test_index = next(sss.split(X, y)) 

X_train, X_test = X[train_index], X[test_index] 
y_train, y_test = y[train_index], y[test_index]

2
文档中提到train_test_split只是在shuffle split基础上的一种方便函数。
我重新排列了一些他们的代码以制作自己的示例。请注意,实际解决方案是代码块的中间部分。其余部分是导入和设置可运行示例。
from sklearn.model_selection import ShuffleSplit
from sklearn.utils import safe_indexing, indexable
from itertools import chain
import numpy as np
X = np.reshape(np.random.randn(20),(10,2)) # 10 training examples
y = np.random.randint(2, size=10) # 10 labels
seed = 1

cv = ShuffleSplit(random_state=seed, test_size=0.25)
arrays = indexable(X, y)
train, test = next(cv.split(X=X))
iterator = list(chain.from_iterable((
    safe_indexing(a, train),
    safe_indexing(a, test),
    train,
    test
    ) for a in arrays)
)
X_train, X_test, train_is, test_is, y_train, y_test, _, _  = iterator

print(X)
print(train_is)
print(X_train)

现在我有实际的索引:train_is, test_is

2
如果您正在使用 pandas,可以通过调用所需模拟数组的 .index 方法来访问索引。 train_test_split 将 pandas 索引传递到新数据框中。
在您的代码中,只需使用 x1.index 返回的数组是与 x 中原始位置相关的索引。

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接