Sklearn将Pandas数据框架和CSR矩阵拆分为测试和训练集

Question

我正在尝试使用scikit learn DecisionTree和Pandas Dataframe对文本进行分类：首先，我构建了一个数据帧，看起来像这样：

   cat1  cat2                             corpus           title
0     0     1                     Test Test Test    erster titel
1     1     0                   Test Super Super   zweiter titel
2     0     1                     Test Test Test   dritter titel
3     0     1                    Test Super Test   vierter titel
4     1     0                   Super Test Super  fuenfter titel
5     1     1         Super einfacher Test Super  fuenfter titel
6     1     1  Super simple einfacher Test Super  fuenfter titel

然后我生成一个TF IDF矩阵：

_matrix = generate_tf_idf_matrix(training_df['corpus'].values)

返回csr矩阵（CountVectorizer->TfidfTransformer）

对于我的分类器，我想使用

    train_X = _matrix
    train_Y = training_df[['cat1','cat2']]

对于多标签分类

我现在的问题是：

如何将我的数据帧和csr矩阵拆分为测试集和训练集？如果我在创建矩阵之前拆分数据帧，csr矩阵会有另一个大小，因为我的文档有不同的功能

限制：我不想把矩阵转换成数组，这样我就可以很容易地分割它。

Answer 1

scikit-learns包已经包含了一个非常强大的模块，用于训练val测试交叉验证功能。您可以快速查看整个模块sklearn.cross_validation（此处为API）。

通常情况下，train_testrongplit可以完成以下工作：

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33)

然而，如果你的类y非常不平衡，你可能会对StratifiedShuffleSplit感兴趣，它在训练/测试数据集中分割数据，但在每个训练/测试集中保留每个类的percentaje。

因此，在您的情况下，首先创建X = _matrix和y = training_df[['cat1', 'cat2']]，然后使用scikit-learn的函数在训练/测试数据集中对其进行拆分。