如何拆分测试和训练数据,以便每个类中至少有一个等级

时间:2015-07-06 05:03:54

标签: python pandas machine-learning scikit-learn classification

我有一些相当不平衡的数据,我试图分类。 但是,它的分类相当不错。

要准确评估效果如何,我必须将数据拆分为训练和测试子集。

现在我通过非常简单的措施来做到这一点:

import numpy as np
corpus = pandas.DataFrame(..., columns=["data","label"]) # My data, simplified
train_index = np.random.rand(len(corpus))>0.2
training_data = corpus[train_index]
test_data = corpus[np.logical_not(train_index)]

这很简单,但有些类非常罕见: 在超过50,000个病例的语料库中,每个约15个发生不到100次,其中两个每个只出现一次。

我想将我的数据语料库划分为测试和训练子集,以便:

  • 如果某个课程的次数少于两次,则会将其排除在
  • 之外
  • 每个班级至少出现一次,在测试和培训中
  • 分成测试和培训是随机的

我可以把一些东西扔到一起, (可能最简单的方法是删除少于2次出现的事物)然后重新采样直到吐出两边都有),但我想知道是否有一个已经存在的干净方法。

我不认为sklearn.cross_validation.train_test_split会为此做,但它存在表明sklearn可能具有此类功能。

1 个答案:

答案 0 :(得分:3)

以下内容符合您将数据划分为测试和培训的3个条件:

#get rid of items with fewer than 2 occurrences.
corpus=corpus[corpus.groupby('label').label.transform(len)>1]

from sklearn.cross_validation import StratifiedShuffleSplit
sss=StratifiedShuffleSplit(corpus['label'].tolist(), 1, test_size=0.5, random_state=None)

train_index, test_index =list(*sss)
training_data=corpus.iloc[train_index]
test_data=corpus.iloc[test_index]

我使用以下虚构数据框测试了上述代码:

#create random data with labels 0 to 39, then add 2 label case and one label case.     
corpus=pd.DataFrame({'data':np.random.randn(49998),'label':np.random.randint(40,size=49998)})
corpus.loc[49998]=[random.random(),40]
corpus.loc[49999]=[random.random(),40]
corpus.loc[50000]=[random.random(),41]

在测试代码时会产生以下输出:

test_data[test_data['label']==40]
Out[110]: 
           data  label
49999  0.231547     40

training_data[training_data['label']==40]
Out[111]: 
           data  label
49998  0.253789     40

test_data[test_data['label']==41]
Out[112]: 
Empty DataFrame
Columns: [data, label]
Index: []

training_data[training_data['label']==41]
Out[113]: 
Empty DataFrame
Columns: [data, label]
Index: []