我有一个csv文件,它被读入数据帧。我根据一列的值将其拆分为训练和测试文件。
让我们说这个专栏被称为"类别"它有几个类别名称作为列值 如cat1,cat2,cat3等不止一次重复。
我需要拆分文件,使每个类别名称至少在这两个文件中都有一次。
到目前为止,我可以根据比例将文件拆分为两个。我尝试了很多选项,但到目前为止这是最好的选择。
def executeSplitData(self):
data = self.readCSV()
df = data
if self.column in data:
train, test = train_test_split(df, stratify = None, test_size=0.5)
self.writeTrainFile(train)
self.writeTestFile(test)
我不完全理解test_train_split中的分层选项。 请帮忙。感谢
答案 0 :(得分:2)
我尝试根据docs使用它,但无法使用stratify
。
from sklearn.cross_validation import train_test_split
import pandas as pd
import numpy as np
np.random.seed([3,1415])
p = np.arange(1, 5.) / np.arange(1, 5.).sum()
df = pd.DataFrame({'category': np.random.choice(('cat1', 'cat2', 'cat3', 'cat4'), (1000,), p=p),
'x': np.random.rand(1000), 'y': np.random.choice(range(2), (1000,))})
def get_freq(s):
return s.value_counts() / len(s)
print get_freq(df.category)
cat4 0.400
cat3 0.284
cat2 0.208
cat1 0.108
Name: category, dtype: float64
如果我尝试:
train, test = train_test_split(df, stratify=df.category, test_size=.5)
train, test = train_test_split(df, stratify=df.category.values, test_size=.5)
train, test = train_test_split(df, stratify=df.category.values.tolist(), test_size=.5)
全部退回了:
TypeError:传递的参数无效:
文档说:
分层:类似数组或无(默认为无)
我想不出为什么这不起作用。
我决定开展一项工作:
def stratify_train_test(df, stratifyby, *args, **kwargs):
train, test = pd.DataFrame(), pd.DataFrame()
gb = df.groupby(stratifyby)
for k in gb.groups:
traink, testk = train_test_split(gb.get_group(k), *args, **kwargs)
train = pd.concat([train, traink])
test = pd.concat([test, testk])
return train, test
train, test = stratify_train_test(df, 'category', test_size=.5)
# this also works
# train, test = stratify_train_test(df, df.category, test_size=.5)
print get_freq(train.category)
print len(train)
Name: category, dtype: float64
cat4 0.400
cat3 0.284
cat2 0.208
cat1 0.108
Name: category, dtype: float64
500
print get_freq(test.category)
print len(test)
cat4 0.400
cat3 0.284
cat2 0.208
cat1 0.108
Name: category, dtype: float64
500