scikit_learn train_test_splitを用いて要素数に偏りがないようにテストデータを作成したい。

Question

前提・実現したいこと

クラスごとのラベルの数に偏りのある学習データをscikit-learnのtrain_test_splitを用いて学習データと訓練データに分割したいと考えています。

発生している問題

クラスごとに、分割する割合を決めたいのですがうまく実装できません。

クラスによってデータを集めることが難しいものを学習データとあるので、どうしても大きな偏りができてしまいます。
そのため、テストデータのクラスごとの要素数が極端に少なくならないようにする方法を探しています。

該当のソースコード

#学習データを想定(1000*2*2*1の配列)
a = np.arange(4000).reshape(1000,2,2,1)

#学習データのラベルを想定(1000*1の配列)
b = np.array([0] * 500 + [1] *200 + [2] * 200 + [3] * 100).reshape(1000,1)

#分割
X_train, X_test, Y_train, Y_test = train_test_split(a, b,test_size=0.1,random_state=1,stratify=b)

#学習データとテストデータの次元と要素数の確認
print(X_train.shape)
print(Y_train.shape)
print(X_test.shape)
print(Y_test.shape)

#学習データとテストデータのクラスの要素数の確認
print((Y_test == 0).sum())
print((Y_test == 1).sum())
print((Y_test == 2).sum())
print((Y_test == 3).sum())

#実行結果
(750, 2, 2, 1)
(750, 1)
(250, 2, 2, 1)
(250, 1)
125
50
50
25

入力データは私が実際に用いるデータの次元を想定し、4次元にしました。

この場合、クラスラベル3の要素数が少ないので、クラス3をもう少し学習データに回したり、テストデータ側に回したり、できればと考えています（例：テストデータ内でクラス3の要素数が25、と少ないので、クラス0 1 2の要素数はそのままに、クラス3だけ50に増やす）。
または、テストデータ内でクラス3の要素数はそのままにして、クラス0 1 2を学習データ側に回して、テストデータ内の各クラスの要素数を統一したいと考えています(例：テストデータ内の要素数を全て50に統一する)。

試したこと

①stratify=に訓練データのラベルの配列を指定した。
②test_size=の値変化させてみた。

調べて見たところ、クラスごとに割合を指定している、文献やサイトは見つけることができませんでした。

補足情報（FW/ツールのバージョンなど）

scilit-learn 0.21.3
python 3.6.8(.pyenv)

どうぞ、よろしくお願い致します。

Comments

[meg_]

一般的にはシャッフルするかと思いますが、それでは不十分なのでしょうか？

[yuki_O_7]

ご返信ありがとうございます。

はい、シャッフルももちろんしたいのですが、データ数が少ないと何かと問題が起きると思うので、柔軟に対応ができるように作成しておきたいと考えております。

[meg_]

機械学習において訓練データのクラスの割合を恣意的に変えることはするものでしょうか？（サンプルの取り方が悪くて極端に偏った場合はサンプル数を増やすべきと思います）
学習したモデルがフィットしているかどうかはテストデータで検証するものではないでしょうか？

Answer 1

現時点でscikit-learnに実装されているもので、その仕様のものは私が確認した限りではありません。念の為、下記を見てみてください（model_selection配下に見つけられなければありません）。

API Reference — scikit-learn 0.22 documentation module-sklearn.model_selection

numpyなどを用いて自分で実装すること自体は可能でしょうから、そうすると良いかと思います。