scikit-learnのtrain_test_splitで、訓練データとテストデータのラベルの割合を揃える

自分の場合なのですが、普段の業務で機械学習を行う場合不均衡データを扱うことが非常に多くあります。
ラベルづけされたデータを train_test_split で訓練データとテストデータに分けるとき、
運が悪いと訓練データとテストデータで、ラベルの割合がずいぶん変わってしまうことがありました。

#　　全データのラベルの割合は　99:1
df['label'].value_counts()
'''
0    9900
1     100
'''
# データの2割りをテストデータにする
df_train, df_test = train_test_split(df, test_size=0.2)

# 訓練データでは ラベル1 は 0.9625 %
df_train.label.value_counts() / len(df_train)
```
0    0.990375
1    0.009625
```
# テストデータでは ラベル1 は 1.15%
df_test.label.value_counts() / len(df_test)
```
0    0.9885
1    0.0115
```

この例ではまだ許容範囲かなという気もしますが運が悪いとかなりの差が開きます。

そこで、かつてはデータフレームをラベルごとに分けてから個別に訓練用とテスト用に分けて、
それをマージして訓練データとテストデータを作ると言った面倒なことをやっていたことがあります。

その後、 train_test_split のマニュアルを読んでいたら非常に便利な引数があることがわかりました。

stratify に、割合を揃えたい列を指定してあげると、訓練データとテストデータで同じ割合になるように分けてくれます。

#　　全データのラベルの割合は　99:1
df['label'].value_counts()
'''
0    9900
1     100
'''
# データの2割をテストデータにする
df_train, df_test = train_test_split(df, test_size=0.2, stratify=df.label)

df_train.label.value_counts() / len(df_train)
```
0    0.99
1    0.01
```
df_test.label.value_counts() / len(df_test)
```
0    0.99
1    0.01
```

綺麗に分かれました。