データフレームの列からapplyで新しい列を作る時、複数列まとめて作成する

DataFrame(その列なので正確にはSeries)に、何か関数を適用して新しい列を作ることは、
機械学習の特徴量作成や前処理などで頻繁に行う処理だと思います。

いつも、1列作るごとに、applyして結果を得ています。
例えば、とある列の値を2乗した列と、3乗した値が欲しいときは次のように書きます。


df = pd.DataFrame(
        np.random.randint(10, size=(5, 3)),
        columns=[f"col{str(i)}" for i in range(3)]
)

# 生成するデータごとにapplyする
df["pow_2"] = df["col0"].apply(lambda x: x**2)
df["pow_3"] = df["col0"].apply(lambda x: x**3)

print(df)

"""
   col0  col1  col2  pow_2  pow_3
0     0     1     6      0      0
1     9     8     8     81    729
2     1     6     8      1      1
3     5     1     2     25    125
4     0     2     5      0      0
"""

このくらい簡単な例であれば、計算負荷も大したことがないのですが、
物によっては、非常に無駄な処理をすることがあります。
例えば日本語の自然言語処理で大量のテキストを形態素解析し、表層形と原形と品詞の列を
それぞれ取得したいときなど、共通の形態素解析処理部分は一回で済ましたいので
3列個別にapplyするなどやりたくありません。

このような場合、applyする関数の戻り値をSeriesで戻せば、
applyの戻りを服す列にできることを知りました。

例えば次のように書きます。


# もう一度サンプルデータ生成
df = pd.DataFrame(
        np.random.randint(10, size=(5, 3)),
        columns=[f"col{str(i)}" for i in range(3)]
)

df[["pow_2", "pow_3"]] = df["col0"].apply(lambda x: pd.Series([x**2, x**3]))
print(df)
"""
   col0  col1  col2  pow_2  pow_3
0     6     3     1     36    216
1     8     4     9     64    512
2     4     4     7     16     64
3     1     1     4      1      1
4     1     8     3      1      1
"""

タプルや配列ではだめで、Seriesで返した場合のみの挙動です。
lambda式を遣わず、普通に定義したSeriesを返す関数でもできます。

なぜこのような挙動になるのか公式ドキュメント内からは該当箇所を探せていないのですが、
とても便利なので積極的に使っていきたいです。

KerasのTokenizerの細かい設定

前回の記事の続きで、 Keras の Tokenizer の話です。

Tokenizerのインスタンスを作成する時、前回紹介した引数の filters / lower/ split /char_level のほか、
num_wordsと、oov_token という二つの引数を設定することができます。
これが少し癖があるので確認しておきましょう。

まず、num_wordsですが、こちらは学習する単語の種類数を指定するものです。
ここで注意しないといけないのは、インデックス0がどの単語にも割り当てられないインデックスとして予約されてしまっているので、
Tokenizerが学習してくれる単語は、num_wordsで指定したのよりも一つ少なくなります。
(普通は数千とかの値を渡すのですが動きを確認しやすくするために、)試しにnum_words=7 を渡してみましょう。
なお、前回の記事同様に、 変数 text_data に、テキストデータが入ってるものとします。


# num_words で、学習する語彙数を指定できる。
keras_tokenizer_7 = Tokenizer(
                            num_words=7
                        )

# 文字列から学習する
keras_tokenizer_7.fit_on_texts(text_data)

# 学習した単語とそのindex
print(keras_tokenizer_7.word_index)
# {'the': 1, 'of': 2, 'to': 3, 'and': 4, 'a': 5, 'in': 6, 'is': 7, 'i': 8, 'that': 9, 'it': 10, 'for': 11, 'this': 12,... 

まず、戸惑うところは、word_indexには大量の単語が含まれている点です。でも、実際にトークン化できる単語は頻度が高い7-1=6個だけになります。
試しに高頻度の単語を並べたテキストをトークン化してみましょう。


print(keras_tokenizer_7.texts_to_sequences(["the of to and a in is i that it for"]))
# [[1, 2, 3, 4, 5, 6]]

ちょっとわかりにくいですが、 the から in までの 6単語が変換され、 is 以降は無視されました。

次に、 oov_token です。 これは元々のテキストに含まれていなかったり、num_wordsで指定した上限値によって溢れてしまった単語を、
oov_token で渡した文字列に変換するものです。
それだけなら単純ですが、注意しないといけないのは、oov_tokenがインデックス1を使ってしまうので、
oov_tokenに何か値を指定すると、学習してくれる単語が一つ減ります。

num_words を先ほどと同じように7にしておくと、 the, of, to, and, a, の5単語しか変換しないことがわかります。


keras_tokenizer_oov = Tokenizer(
                            num_words=7,
                            oov_token="未知語"
                        )

# 文字列から学習する
keras_tokenizer_oov.fit_on_texts(text_data)

# 学習した単語とそのindex
print(keras_tokenizer_oov.word_index)
# {'未知語': 1, 'the': 2, 'of': 3, 'to': 4, 'and': 5, 'a': 6, 'in': 7, 'is': 8, 'i': 9, 'that': 10, 'it': 11,...

print(keras_tokenizer_oov.texts_to_sequences(["the of to and a in is i that it for"]))
# [[2, 3, 4, 5, 6, 1, 1, 1, 1, 1, 1]]

oov_token を使うと、元の単語列と、変換後の数列の要素数が維持されますし、
未学習だった単語がどれだけ含まれているのかもわかるのでとても便利ですが、
使える単語が一個減ることなどは注意しないといけませんね。

ちなみに、 oov_token に、 学習対象のドキュメントにある単語を使うと、その単語のインデックスがoov_tokenとして使われ、
1が使われなくなります。
ちょっとややこしいのでできるだけ避けましょう。

KerasのTokenizerの基本的な使い方

自然言語処理において翻訳などのseq2seqモデルやそれ以外でもRNN系のモデルを使う場合、
前処理においてテキストの列を数列に変換(トークン化)することがあります。

そのよな時に、Kerasのユーティリティーに用意されている、Tokenizerが便利なのでその基本的な使い方を紹介します。
今回の主な内容は次の4つです。(その他細かいオプションとか、別の使い側は次回以降の更新で。)
– インスタンスの生成
– テキストを数列化する
– デフォルトパラメーターで生成した時の設定
– 数列をテキストに戻す

ドキュメントはこちらです。

サンプルに何かデータが必要なので、20newsのデータを一部だけ読み込んで使います。


from sklearn.datasets import fetch_20newsgroups

# データの読み込み。少量で良いのでカテゴリも一つに絞る。
remove = ('headers', 'footers', 'quotes')
categorys = [
        "sci.med",
    ]
twenty_news = fetch_20newsgroups(
                                subset='train',
                                remove=remove,
                                categories=categorys
                            )
text_data = twenty_news.data

Tokenizer を使うときはまずはインスタンスを生成し、
テキストデータを学習させる必要があります。
(ここで学習しなかった単語はトークン化できません。)


from tensorflow.keras.preprocessing.text import Tokenizer
# Tokenizer のインスタンス生成
keras_tokenizer = Tokenizer()
# 文字列から学習する
keras_tokenizer.fit_on_texts(text_data)

# 学習した単語とそのindex
print(keras_tokenizer.word_index)
"""
{'the': 1, 'of': 2, 'to': 3, 'and': 4, 'a': 5, 'in': 6, 'is': 7,
 'i': 8, 'that': 9, 'it': 10, 'for': 11, 'this': 12, 'are': 13, ...,
--- 以下略 ---
"""

テキストデータをトークン化するときは、texts_to_sequences に、”テキストデータの配列を”渡します。
テキストを一つだけ渡すと、それを文字単位に分解してしまうので注意してください。


# テキストデータを数列に変更
sequence_data = keras_tokenizer.texts_to_sequences(text_data)
# 一つ目のテキストの変換結果。
print(sequence_data[0])
"""
[780, 3, 1800, 4784, 4785, 3063, 1800, 2596, 10, 41, 130, 24,
15, 4, 148, 388, 2597, 11, 60, 110, 20, 38, 515, 108, 586, 704,
353, 21, 46, 31, 7, 467, 3, 268, 6, 5, 4786, 965, 2223, 43, 2598,
2, 1, 515, 24, 15, 13, 747, 11, 5, 705, 662, 586, 37, 423, 587, 7092,
77, 13, 1490, 3, 130, 16, 5, 2224, 2, 12, 415, 3064, 12, 7, 5, 516,
6, 40, 79, 47, 18, 610, 3732, 1801, 26, 2225, 706, 918, 3065, 2,
1, 1801, 21, 32, 61, 1638, 31, 329, 7, 9, 1, 1802, 966, 1491,
18, 3, 126, 3066, 4, 50, 1352, 3067]
"""

これで目的のトークン化ができました。
今回は、インスタンス化する時に何も引数を渡さず、完全にデフォルトの設定になっているのですが、
一応主な設定を確認しておきましょう。


# デフォルトでは、文字を小文字に揃える。
print(keras_tokenizer.lower)
# True

# デフォルトでは文字単位ではなく、次のsplitで区切った単語単位でトークン化する。
print(keras_tokenizer.char_level)
# False

# デフォルトでは、split に半角スペースが指定されており、スーペースで区切られる。
print(keras_tokenizer.split == " ")

# いくつかの記号は除外され、単語中に含まれている場合はそこで区切られる。
print(keras_tokenizer.filters)
# !"#$%&()*+,-./:;<=>?@[\]^_`{|}~

# 例えば、 dog&cat は &が取り除かれ、 dog と cat が個別にトークン化される。
print(keras_tokenizer.texts_to_sequences(["dog&cat"]))
# [[7316, 2043]]

# & が 半角ペースだった場合と結果は同じ
print(keras_tokenizer.texts_to_sequences(["dog cat"]))
# [[7316, 2043]]

最後に、トークン列をテキストに戻す方法です。
sequences_to_texts を使います。


# 数列をテキストに戻す。
text_data_2 = keras_tokenizer.sequences_to_texts(sequence_data)

print("元のテキスト")
print(text_data[0])
print("\n復元したテキスト")
print(text_data_2[0])

"""
元のテキスト
[reply to keith@actrix.gen.nz (Keith Stewart)]


It would help if you (and anyone else asking for medical information on
some subject) could ask specific questions, as no one is likely to type
in a textbook chapter covering all aspects of the subject.  If you are
looking for a comprehensive review, ask your local hospital librarian.
Most are happy to help with a request of this sort.

Briefly, this is a condition in which patients who have significant
residual weakness from childhood polio notice progression of the
weakness as they get older.  One theory is that the remaining motor
neurons have to work harder and so die sooner.

復元したテキスト
reply to keith actrix gen nz keith stewart it would help if you and anyone
else asking for medical information on some subject could ask specific
questions as no one is likely to type in a textbook chapter covering all
aspects of the subject if you are looking for a comprehensive review ask
your local hospital librarian most are happy to help with a request of this
sort briefly this is a condition in which patients who have significant
residual weakness from childhood polio notice progression of
the weakness as they get older one theory is that the remaining
motor neurons have to work harder and so die sooner
"""

改行のほか、括弧やカンマなどの記号が消えていること、一部の大文字が小文字になっていることなどが確認できます。

kerasのモデルの可視化

kerasでモデルを構築したとき、構築したモデルが意図した構造になっているかどうか可視化して確認する方法です。
Sequentialモデルであれば、 .summary()で十分なことが多いのですが、functional APIを使って複雑なモデルを作る場合に重宝します。

kerasのドキュメントを見ると、そのままズバリな名前で 可視化 のページがあり、plot_modelという関数が説明されています。
可視化 – Keras Documentation

「graphvizを用いて」と書かれている通り、graphvizがインストールされている必要がありますが、
このほか pydot というライブラリも必要なのでpip等でインストールしておきましょう。
(他サイトなどでpydotは開発が止まっていて動かないからpydotplusを使う、といった趣旨の記事を見かけますが、
現在はpydotの開発が再開されているようでpydotで動きます。)

さて、graphvizとpydotが入ったら、早速ちょっとだけ複雑なモデルを作ってみて、可視化してみましょう。
一応 model.summary() の結果も表示してみました。

まずは可視化対象のモデル構築から。


from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Concatenate
from tensorflow.keras.layers import Add

i0 = Input(shape=(64, ))
i1 = Input(shape=(64, ))
x0 = Concatenate()([i0, i1])
x1 = Dense(32, activation="tanh")(x0)
x2 = Dense(32, activation="tanh")(x1)
x3 = Add()([x1, x2])
x4 = Dense(1, activation="sigmoid")(x3)
model = Model([i0, i1], x4)

print(model.summary())
# 以下出力結果
"""
Model: "model"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            [(None, 64)]         0                                            
__________________________________________________________________________________________________
input_2 (InputLayer)            [(None, 64)]         0                                            
__________________________________________________________________________________________________
concatenate (Concatenate)       (None, 128)          0           input_1[0][0]                    
                                                                 input_2[0][0]                    
__________________________________________________________________________________________________
dense (Dense)                   (None, 32)           4128        concatenate[0][0]                
__________________________________________________________________________________________________
dense_1 (Dense)                 (None, 32)           1056        dense[0][0]                      
__________________________________________________________________________________________________
add (Add)                       (None, 32)           0           dense[0][0]                      
                                                                 dense_1[0][0]                    
__________________________________________________________________________________________________
dense_2 (Dense)                 (None, 1)            33          add[0][0]                        
==================================================================================================
Total params: 5,217
Trainable params: 5,217
Non-trainable params: 0
__________________________________________________________________________________________________
None
"""

Connected to に複数レイヤー入っているとぱっと見わかりにくいですね。

次にplot_model使ってみます。
show_shapes オプションを使って、入出力の形も表示してみました。


from tensorflow.keras.utils import plot_model
plot_model(
    model,
    show_shapes=True,
)

出力されたのがこちら。

モデルの形をイメージしやすいですね。

matplotlibでgif動画生成

3次元グラフの次は動画(gif)を用いたデータの可視化方法のメモです。

matplotlibでは、次のクラスに、パラパラ漫画のようにグラフのリストを渡してあげることで、
アニメーションさせることができます。
matplotlib.animation.ArtistAnimation

例として、サインカーブを少しずつずらしながら描いてみました。


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import ArtistAnimation

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)

# 0 <=x < 2pi の範囲の点列を作成。
x = np.linspace(0, 2*np.pi, 101)[: -1]
# 各コマの画像を格納する配列
image_list = []

for i in range(100):
    # ずらしながらsinカーブを描写し、配列に格納
    y = np.sin(np.roll(x, -i))
    image = ax.plot(x, y)
    image_list.append(image)

# アニメーションを作成
ani = ArtistAnimation(fig, image_list, interval=10)
# gifに保存
ani.save('sin_animation.gif', writer='pillow')

保存されたgifがこちらです。

フルサイズで貼り付けると記事中でも動くのですね。
(いつものようにサムネイルで張ったら止まってしまっていて、クリックしないと動画になりませんでした。)

動画が使えると少しデータの可視化の幅が広がりそうです。
とりあえず機械学習の学習の進捗とかの可視化などに使ってみたいです。

matplotlibでSurface plots

昨日に続いてmatplotlibの3次元グラフの話です。
今回のテーマは Surface plots。(日本語では表面プロットでいいのかな?)
2変数関数の可視化等に便利なやつですね。

ドキュメントは今回もこちら。 : The mplot3d Toolkit

今回は例として 鞍点を持つ次の関数を可視化してみましょう。
$$
z = f(x, y) = x^2 – y^2.
$$

比較用に等高線で可視化したグラフと並べてみました。
参考: matplotlibで等高線


import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np


def f(x, y):
    return x**2 - y**2


X, Y = np.meshgrid(
        np.linspace(-10, 10, 101),
        np.linspace(-10, 10, 101),
    )
Z = f(X, Y)

fig = plt.figure(figsize=(16, 6), facecolor="w")
ax_3d = fig.add_subplot(121, projection="3d")
ax_3d.plot_surface(X, Y, Z)

ax = fig.add_subplot(122)
contour = ax.contourf(X, Y, Z)
fig.colorbar(contour)
plt.show()

結果はこちら。

可視化する対象によって向き不向きがあるのでいつもそうだというわけではないのですが、
今回のサンプルでは圧倒的に3次元プロットの方が圧倒的に関数の形をつかみやすいですね。

matplotlibで3D散布図

matplotlibで3次元のグラフを作成する方法のメモです。
今回は散布図を描いてみます。

matplotlibで3次元のグラフを書くには、mplot3d Toolkitというのを使います。
ドキュメント: The mplot3d Toolkit
また、 3次元散布図についてはこちらのドキュメントも参考になります。 3D scatterplot

ポイントとしては、(importした後明示的には使わないので忘れがちですが、)
Axes3Dを必ずインポートしておくことと、axを取得するときに、
projection="3d"を忘れないことですね。

iris のデータの4つの特徴量の中から適当に3個選んでやってみます。


import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn.datasets import load_iris

iris = load_iris()
data = iris.data
label = iris.target

fig = plt.figure(figsize=(8, 8), facecolor="w")
ax = fig.add_subplot(111, projection="3d")
for c in range(3):
    ax.scatter(
        data[label == c, 0],
        data[label == c, 2],
        data[label == c, 3],
        label=iris.target_names[c]
    )

ax.set_xlabel(iris.feature_names[0])
ax.set_ylabel(iris.feature_names[2])
ax.set_zlabel(iris.feature_names[3])

ax.legend()
plt.show()

結果がこちら。
綺麗に3次元のプロットができました。

pandasのSeriesを辞書型に変換する3つの方法

以前書いた、 DataFrameの2列の値からdictを作る に近い話です。
今回の対象は Dataframeではなく Series。
元々、辞書と同じようにSeries[kye]で、値を取り出すことができるので、
DataFrameに比べて辞書型に変換するニーズも少ないと思うのですが、
to_dict() メソッドを持ってることを最近知ったのでこの記事を書きました。

まず、サンプルルとなるデータを作ります。


import pandas as pd
data = pd.Series({chr(i): i for i in range(97, 105)})
print(data)
"""
a     97
b     98
c     99
d    100
e    101
f    102
g    103
h    104
dtype: int64
"""

それでは、これを3種類の方法で辞書に変換してみましょう。
それぞれ、内包表記を使う方法(昔よく使っていた。)、 to_dict()を使う方法、 dict()にいれてしまう方法(一番楽)です。


print({k: v for k, v in chr_codes.items()})
# {'a': 97, 'b': 98, 'c': 99, 'd': 100, 'e': 101, 'f': 102, 'g': 103, 'h': 104}

print(chr_codes.to_dict())
# {'a': 97, 'b': 98, 'c': 99, 'd': 100, 'e': 101, 'f': 102, 'g': 103, 'h': 104}

print(dict(chr_codes))
# {'a': 97, 'b': 98, 'c': 99, 'd': 100, 'e': 101, 'f': 102, 'g': 103, 'h': 104}

dictに変換するメリットとして、dictならgetメソッドが使えて、keyの中に取りたい値がない場合のデフォルト値の設定もできますよっていう話を
最後に書こうと思っていたのですが、
Seriesにも getメソッドは用意されていて、ほぼ同じように動作するようでした。
pandas.Series.get

どちらかというと、辞書に変換する方法よりも辞書に変換しなくても
同じように使えるってことを覚えておいた方が有益に思います。

matplotlibで二重軸

matplotlibで一つの枠の中に複数のグラフを書く場合、値のレンジが近ければ良いのですが、
10倍以上も違うと片方のグラフが潰れてしまうなど、不便なことがあります。

そのような時はだいたいグラフを分けて書いたり、
Tableauなどの別のソフトを使って2重軸のグラフを書くなどの対応をしているのですが、
matplotlibでも左右の軸を使ったグラフを書くことはできます。

その際は、 twinx というメソッドを使います。
以下、サンプルコードです。
この時、凡例をつけておかないと、それぞれのグラフがどちらの軸を見るのかわからないので、つけるのですが、
ちょっとつけかたが特殊なので、そのサンプルも兼ねています。
(get_legend_handles_labels というメソッドを使います。)


import matplotlib.pyplot as plt
import numpy as np

# データの作成
x = np.arange(0, 20)
y1 = np.random.randint(300, 500, size=20)
y2 = np.random.randint(0, 20, size=20)

fig = plt.figure(facecolor="w")
ax1 = fig.add_subplot(1, 1, 1)
ax1.plot(x, y1, label="折れ線グラフ")
ax1.set_ylabel("折れ線グラフ")
ax1.set_ylim([200, 500])

# 2重軸の作成
ax2 = ax1.twinx()
ax2.bar(x, y2, label="棒グラフ", color="g", alpha=0.5)
ax2.set_ylabel("棒グラフ")
ax2.set_ylim([0, 60])

# 凡例をまとめて出力する
handler1, label1 = ax1.get_legend_handles_labels()
handler2, label2 = ax2.get_legend_handles_labels()
ax1.legend(handler1 + handler2, label1 + label2)

plt.show()

出力される図がこちら。

層化K分割交差検証の紹介とPythonで実行する方法

少し前の記事になりますが、 scikit-learnでK-分割交差検証 というのを書きました。
これは、分類のタスクでは目的変数の件数がクラスごとにある程度揃っていたり、データが十分に揃っていればうまく機能します。
しかし、一方で不均衡データなど、目的変数の値の割合が偏っていて特に、一部のクラスのデータが非常に少ないと困ったことになります。

試しに、いつものirisのデータを少し絞り込んで、元々種類ごとに50件ずつあるデータを
setosa: 50個
versicolor: 10個
virginica: 5個
にして試してみます。
(一番少ないクラスのデータ件数が5個なのに5分割するという極端な例ですが、
説明のためなのでご了承ください。)


from sklearn.model_selection import KFold
from sklearn.datasets import load_iris
import numpy as np

# データの読み込み
iris = load_iris()
X = iris.data
y = iris.target

# 実験のため対象を絞り込んで不均衡データにする
index = list(range(50)) + list(range(50, 60)) + list(range(100, 105))
X = X[index]
y = y[index]

for c in range(3):
    print(f"{iris.target_names[c]}: {list(y).count(c)}個")

"""
setosa: 50個
versicolor: 10個
virginica: 5個
"""

# KFoldを用いてK-分割交差検証した時に各グループに含まれるラベル数
kf = KFold(5, shuffle=True)
i = 0
for train_index, test_index in kf.split(X):
    i += 1
    print(f"\n{i}グループの訓練データに含まれるラベル")
    train_y = y[train_index]
    for c in range(3):
        print(f"{iris.target_names[c]}: {list(train_y).count(c)}個")
    print(f"{i}グループのテストデータに含まれるラベル")
    test_y = y[test_index]
    for c in range(3):
        print(f"{iris.target_names[c]}: {list(test_y).count(c)}個")

"""
1グループの訓練データに含まれるラベル
setosa: 41個
versicolor: 6個
virginica: 5個
1グループのテストデータに含まれるラベル
setosa: 9個
versicolor: 4個
virginica: 0個

2グループの訓練データに含まれるラベル
setosa: 42個
versicolor: 7個
virginica: 3個
2グループのテストデータに含まれるラベル
setosa: 8個
versicolor: 3個
virginica: 2個

3グループの訓練データに含まれるラベル
setosa: 39個
versicolor: 9個
virginica: 4個
3グループのテストデータに含まれるラベル
setosa: 11個
versicolor: 1個
virginica: 1個

4グループの訓練データに含まれるラベル
setosa: 40個
versicolor: 8個
virginica: 4個
4グループのテストデータに含まれるラベル
setosa: 10個
versicolor: 2個
virginica: 1個

5グループの訓練データに含まれるラベル
setosa: 38個
versicolor: 10個
virginica: 4個
5グループのテストデータに含まれるラベル
setosa: 12個
versicolor: 0個
virginica: 1個
"""

結果が長くなって恐縮ですが、1グループ目では、テストデータにvirginicaが含まれなくなっていますし、
5グループ目では versicolor がテストデータにありません。
逆に、訓練データでそれらのデータの割合が過剰に高くなっています。
これではモデルの学習もうまくいきませんし、評価も適切に行えません。

このような時、 train_test_splitであれば、stratify引数を使って、ラベルの割合を揃えられます。
参考: scikit-learnのtrain_test_splitで、訓練データとテストデータのラベルの割合を揃える
そして、KFoldには stratify がないのですが代わりに、
タイトルの 層化K分割交差検証(Stratified K-Folds cross-validator)という手法が知られており、それに対応する
StratifiedKFold というクラスが用意されています。

要は、ラベルの割合を揃えながらK分割交差検証する方法です。
使い方はKFoldととても似ていますが、splitするときに、labelも渡してあげる必要がある点だけ注意です。


from sklearn.model_selection import StratifiedKFold

skf = StratifiedKFold(5, shuffle=True)
i = 0
for train_index, test_index in skf.split(X, y):
    i += 1
    print(f"\n{i}グループの訓練データに含まれるラベル")
    train_y = y[train_index]
    for c in range(3):
        print(f"{iris.target_names[c]}: {list(train_y).count(c)}個")
    print(f"{i}グループのテストデータに含まれるラベル")
    test_y = y[test_index]
    for c in range(3):
        print(f"{iris.target_names[c]}: {list(test_y).count(c)}個")

"""
1グループの訓練データに含まれるラベル
setosa: 40個
versicolor: 8個
virginica: 4個
1グループのテストデータに含まれるラベル
setosa: 10個
versicolor: 2個
virginica: 1個

2グループの訓練データに含まれるラベル
setosa: 40個
versicolor: 8個
virginica: 4個
2グループのテストデータに含まれるラベル
setosa: 10個
versicolor: 2個
virginica: 1個

3グループの訓練データに含まれるラベル
setosa: 40個
versicolor: 8個
virginica: 4個
3グループのテストデータに含まれるラベル
setosa: 10個
versicolor: 2個
virginica: 1個

4グループの訓練データに含まれるラベル
setosa: 40個
versicolor: 8個
virginica: 4個
4グループのテストデータに含まれるラベル
setosa: 10個
versicolor: 2個
virginica: 1個

5グループの訓練データに含まれるラベル
setosa: 40個
versicolor: 8個
virginica: 4個
5グループのテストデータに含まれるラベル
setosa: 10個
versicolor: 2個
virginica: 1個
"""

全グループで、訓練データとテストデータの割合が揃っているのを確認できました。