在 scikit 中使用训练测试拆分时获取索引答案

【问题标题】：Getting indices while using train test split in scikit在 scikit 中使用训练测试拆分时获取索引
【发布时间】：2016-06-07 22:30:27
【问题描述】：

为了将我的数据分别拆分为训练和测试数据，我正在使用

sklearn.cross_validation.train_test_split 函数。

当我将数据和标签作为列表列表提供给此函数时，它会在两个单独的列表中返回训练和测试数据。

我想从原始数据列表中获取训练和测试数据元素的索引。

谁能帮我解决这个问题？

提前致谢

【问题讨论】：

也回答了here。

标签： python-2.7 scikit-learn

【解决方案1】：

您可以提供索引向量作为附加参数。使用来自sklearn的示例：

import numpy as np
from sklearn.cross_validation import train_test_split
X, y,indices = (0.1*np.arange(10)).reshape((5, 2)),range(10,15),range(5)
X_train, X_test, y_train, y_test,indices_train,indices_test = train_test_split(X, y,indices, test_size=0.33, random_state=42)
indices_train,indices_test
#([2, 0, 3], [1, 4])

【讨论】：

非常简单实用的答案，太棒了！

【解决方案2】：

尝试以下解决方案（取决于您是否有不平衡）：

NUM_ROWS = train.shape[0]
TEST_SIZE = 0.3
indices = np.arange(NUM_ROWS)

# usual train-val split
train_idx, val_idx = train_test_split(indices, test_size=TEST_SIZE, train_size=None)

# stratified train-val split as per Response's proportion (if imbalance)
strat_train_idx, strat_val_idx = train_test_split(indices, test_size=TEST_SIZE, stratify=y)

【讨论】：