使用sklearn分割数据以进行随机训练和测试

时间:2019-10-21 13:17:50

标签: python-3.x scikit-learn

我有一个数据文件,例如:

示例:

     X  Y month  day  FFMC    DMC     DC   ISI  RH  wind  rain   area
68   2  2   sep  fri  92.4  117.9  668.0  12.2  33   6.3   0.0   0.00
228  4  6   sep  sun  93.5  149.3  728.6   8.1  26   3.1   0.0  64.10
387  5  5   mar  thu  90.9   18.9   30.6   8.0  48   5.4   0.0   0.00

我正在尝试将其随机拆分为训练集和测试集,但要基于列而不是行,从第3列到末尾,并且训练集和测试集都将包含前2列。为此,我尝试使用:

from sklearn.cross_validation import train_test_split
data = pd.read_csv('mydata.txt', sep="\t")
data_train, data_test = train_test_split(data, test_size=0.3)

但是此程序包将行而不是列分开。然后我尝试转置文件并使用相同的包,如下所示:

X_train, X_test = train_test_split(data.T, test_size=0.3)

这是预期的输出:

火车组:

     X  Y   month   day FFMC    DC  ISI RH  area
68   2  2   sep fri 92.4    668.0   12.2    33  0.00
228  4  6   sep sun 93.5    728.6   8.1 26  64.10
387  5  5   mar thu 90.9    30.6    8.0 48  0.00

测试集:

     X  Y   DMC wind    rain
68   2  2   117.9   6.3 0.0
228  4  6   149.3   3.1 0.0
387  5  5   18.9    5.4 0.0

您知道我如何修复代码以获得预期的训练和测试集吗?

1 个答案:

答案 0 :(得分:0)

您可以通过编写自己的函数来手动完成此操作:

import pandas as pd
import numpy as np

data = pd.read_csv('mydata.txt', sep="\t")

columns = data.columns
keep = ['X','Y']

def split_columns(columns_list, keep_columns, frac=0.2):

    # remove the common columns for the moment
    for col in keep_columns:
        columns.remove(col)

    # shuffle the rest of the column list
    np.random.shuffle(columns)

    # select the right proportion for the train and test set
    cut = max(1, int((1-frac)*len(columns)))

    train = columns[:cut]
    test = columns[cut:]

    # Add common columns to both lists 
    train.extend(keep_columns)
    test.extend(keep_columns)

    return train, test

train_columns, test_columns = split_columns(columns, keep)

# Build your train and test set by selecting the appropriate subset of columns
train = data[train_columns]
test = data[test_columns]