sklearn:如何将测试数据与原始数据相关联?

时间:2015-08-11 01:43:49

标签: python scikit-learn

我使用sklearn train_test_split函数来分割我的训练和测试数据。在分割数据并运行我的分类器之后,我需要能够将功能和标签值追溯到原始数据记录。我怎样才能做到这一点?有没有办法包含某种被分类器忽略的隐藏ID功能?

import json
import numpy as np
from sklearn.cross_validation import train_test_split

json_data = r"""
[
    { "id": 101, "label": 1, "f1": 1, "f2":2, "f3": 3 },
    { "id": 653, "label": 0, "f1": 2, "f2":7, "f3": 8 },
    { "id": 219, "label": 0, "f1": 4, "f2":9, "f3": 2 },
    { "id": 726, "label": 1, "f1": 6, "f2":1, "f3": 0 },
    { "id": 403, "label": 0, "f1": 1, "f2":5, "f3": 4 }
]"""

data = json.loads(json_data)

feature_names = ["f1", "f2", "f3"]

labels = []
features = []
for item in data:
    temp_list = []
    labels.append(item["label"])
    for feature_name in feature_names:
        temp_list.append(item[feature_name])
    features.append(temp_list)

labels_train, labels_test, features_train, features_test = train_test_split(labels, features, test_size = .20, random_state = 99)

print labels_test
print features_test

## this will give us labels_test = [0], features_test = [[4,9,2]] which corresponds to record with id = 219
## how can I efficiently correlate the split data back to the original records without comparing feature values?

1 个答案:

答案 0 :(得分:1)

通常我将输入数据存储在Pandas数据框中,并使用索引进行训练测试分割;对于你的例子,你可以使用这样的东西:

import pandas as pd
test_size=0.2

df = pd.read_json(json_data)
I = np.random.random(len(df)) > test_size

X_train = df[I][feature_names].values
X_test = df[~I][feature_names].values
y_train = df[I]['label'].values
y_test = df[~I]['label'].values