带有索引的Scikit学习train_test_split

Question

使用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，然后再次扩展。有没有更清洁的解决方案？

Answer 1

Scikit learning在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