train_test_splitの引数について詳しく調べてみた

Taku

kaggleなどのコンペで、学習データと評価用データに分ける際に使うtrain_test_split。引数が多く、整理のためにここにまとめておく。

一般的にデータを分割する際には以下のようなコードを記述すると思う。

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=1)

ではこのコードの引数の意味について詳しくみていこう。

上記リンクのsklearnの公式ドキュメントを参考にした。

・test_size
テストサイズと訓練サイズをそれぞれ何:何に分けるかを定義する（今回は0.2なので、8:2に分割。一般的には0.3や0.2くらいに設定する。）

・random_state
数値を指定することにより、決められたrandom_seedで振り分けてくれる。
誰でも同じ結果を得られるようにすることが目的。

・shuffle
分類する前にデータをシャッフルして分割してくれる。ただし、stratifyの引数を指定する場合は必ずNone（デフォルト）のままにしておく必要がある。

・stratify
データを分割した際、偏りが生じることがある。
（2値分類の場合に、正例の数が少ないため、訓練用データにばかりに正例入って、テストデータにはほとんど入っていないことが起こりうる。この引数を設定することで、正例と負例の割合を訓練データと評価データで揃えることができる。


今回はtrain_test_splitについてまとめてみた。test_sizeの指定はモデルの精度にも関わってくるので色々な比率で試してみると良い。sklearnのホームページには他にも色々なライブラリについての詳しい説明があるので、ぜひ参考にしてほしい。