时间戳

首页 > 解决方案 > 使用 sklearn 拆分数据以随机训练和测试

python-3.x - 使用 sklearn 拆分数据以随机训练和测试

问题描述

我有一个像这个例子的数据文件：

例子：

     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

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

标签： python-3.xscikit-learn

解决方案

您可以通过编写自己的函数来手动完成：

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]

推荐阅读