sklearnのtrain_test_splitを使うときはstratifyを指定した方が良い
はじめに
train_test_splitはsklearnをはじめて学んだ頃からよくお世話になっています。しかし、stratifyを指定しないとまずいことが起こり得ると最近気づきました。
stratifyって何？
層化という言葉を聞いたことがある方が一定数いると思いますが、それです。あるいは、交差検証でStratifiedKFoldを使ったことのある人もだいたい理解しているでしょう。
要するに、クラスラベル（など）ごとにサンプルを取ってくるということを意味します。2クラス分類、100サンプルで元の各クラスの比率が50,50であれば、10件取り出しても5,5ずつになることが保証される、というのがここでいうstratify（層化）の意味です。
これを指定しないと、けっこういい加減なことになります。
指定しないで試してみる
まずdigitsでやります。こんな感じ。
from sklearn.datasets import load_digits from sklearn.model_selection import train_test_split from sklearn.svm import SVC from sklearn.metrics import classification_report def main(): dataset = load_digits() X_train, X_test, y_train, y_test\ = train_test_split(dataset.data, dataset.target) svm = SVC(gamma="scale") svm.fit(X_train, y_train) prediction = svm.predict(X_test) print(classification_report(y_test, prediction)) if __name__ == "__main__": main()
結果
precision recall f1-score support
0 1.00 1.00 1.00 58
1 0.94 1.00 0.97 49
2 1.00 1.00 1.00 46
3 1.00 1.00 1.00 51
4 1.00 0.98 0.99 42
5 1.00 0.98 0.99 42
6 0.98 1.00 0.99 43
7 1.00 1.00 1.00 46
8 0.98 0.93 0.96 46
9 1.00 1.00 1.00 27
accuracy 0.99 450
macro avg 0.99 0.99 0.99 450
weighted avg 0.99 0.99 0.99 450
supportに注目してください。0は58, 9は27で、倍以上ばらついています。あまり良い評価とは言えないということです。
stratifyを指定します。
X_train, X_test, y_train, y_test\
= train_test_split(dataset.data, dataset.target,
stratify=dataset.target)
結果
precision recall f1-score support
0 1.00 0.98 0.99 45
1 0.96 1.00 0.98 46
2 1.00 1.00 1.00 44
3 1.00 1.00 1.00 46
4 0.98 0.96 0.97 45
5 1.00 1.00 1.00 46
6 1.00 1.00 1.00 45
7 1.00 0.98 0.99 45
8 0.98 0.95 0.96 43
9 0.96 1.00 0.98 45
accuracy 0.99 450
macro avg 0.99 0.99 0.99 450
weighted avg 0.99 0.99 0.99 450
若干のぶれがありますが、概ね数がそろっています。これによってより妥当な評価が可能になると思われます。
次に不均衡データの場合について考えてみます。同じデータセットで、5かどうかを判別するとしましょう。
from sklearn.datasets import load_digits from sklearn.model_selection import train_test_split from sklearn.svm import SVC from sklearn.metrics import classification_report def main(): dataset = load_digits() y = dataset.target == 5 X_train, X_test, y_train, y_test\ = train_test_split(dataset.data, y) svm = SVC(gamma="scale") svm.fit(X_train, y_train) prediction = svm.predict(X_test) print(classification_report( y_test, prediction)) if __name__ == "__main__": main()
結果（何回かやった中で偏りがひどかった奴です）
precision recall f1-score support
False 1.00 1.00 1.00 415
True 1.00 0.97 0.99 35
accuracy 1.00 450
macro avg 1.00 0.99 0.99 450
weighted avg 1.00 1.00 1.00 450
stratifyを指定します。
X_train, X_test, y_train, y_test\
= train_test_split(dataset.data, y, stratify=y)
結果
precision recall f1-score support
False 1.00 1.00 1.00 404
True 0.98 0.96 0.97 46
accuracy 0.99 450
macro avg 0.99 0.98 0.98 450
weighted avg 0.99 0.99 0.99 450
この場合は、何回やってもデータ個数の比率は保証されます。
まとめ
ということで、stratifyを指定した方が良いでしょう。
testに回すサンプル数が少ないときは特に相対的にクラス間の比率のばらつきが大きくなるので、やっておくべきだと思います。testに回すサンプル数が多い場合はまだやらなくてもなんとかなるかもしれませんが、それでもやった方が結果が安定するはずです。
