Scikit-learnでクラス数を考慮してデータを分割したい


以下のように、scikit-learnの関数train_test_splitを用いると、

[code]

    dataset_train, dataset_test = train_test_split(dataset, train_size=0.8)
    
[/code]

データセットを訓練データとテストデータに分割はしてくれるのですが、  
クラス数が多い（例えば100クラス）場合だと、  
訓練データとテストデータの各々のクラス数が異なる時があります。  
例えば訓練データに含まれるクラス数は100である一方、  
テストデータのそれは98となってしまうことがあります。

train_test_splitでは、ランダムシャッフルしてsplitしているだけなので、  
クラスに含まれるデータ数がアンバランスな場合、このようなことが起きると思われます。

クラス数もちゃんと保つように、データを分割するには、  
どのようにすれば良いでしょうか？

よろしくお願いします。




クラス比例を保ちながら分割するには`StratifiedShuffleSplit`を[使用して下さい](http://scikit-
learn.org/stable/modules/generated/sklearn.model_selection.StratifiedShuffleSplit.html#sklearn.model_selection.StratifiedShuffleSplit)。

