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

Question

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

sklearn.cross_validation.train_test_split功能。

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

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

谁能帮我解决这个问题？

提前致谢

Answer 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])

Answer 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)