Python: pandas の DataFrame を scikit-learn で KFold するときの注意点

今回は pandas の DataFrame を scikitl-learn で交差検証しようとしてハマった話について。だいぶ平凡なミスなんだけど、またやるとこわいので自分用にメモしておく。

使った環境は次の通り。

$ sw_vers 
ProductName:    Mac OS X
ProductVersion: 10.13.5
BuildVersion:   17F77
$ python -V                         
Python 3.6.5

下準備

まずは pandas と scikit-learn をインストールしておく。

$ pip install pandas scikit-learn scipy

Python のインタプリタを起動する。

$ python

なんか適当に DataFrame を作っておく。

>>> import pandas as pd
>>> data = [
...     ('Ant'),
...     ('Beetle'),
...     ('Cat'),
...     ('Deer'),
...     ('Eagle'),
...     ('Fox'),
... ]
>>> columns = ['name']
>>> df = pd.DataFrame(data, columns=columns)

できた DataFrame はこんな感じ。

>>> df
     name
0     Ant
1  Beetle
2     Cat
3    Deer
4   Eagle
5     Fox

特に指定しない限り、デフォルトではインデックスとして 0 から始まる整数が連番で振られる。

sklearn.model_selection.KFold の使い方

scikit-learn で交差検証するとき基本は KFold クラスを使う。このクラスはインスタンス化するときに分割数を指定し、その上で KFold#split() メソッドに分割するものを渡す。返り値としてはイテラブルなオブジェクトが返ってきて、それぞれ学習用データと検証用データ用のインデックスが取り出せる。

>>> from sklearn.model_selection import KFold
>>> kf = KFold(n_splits=3)
>>> for train, test in kf.split(df):
...     print(train, test)
... 
[2 3 4 5] [0 1]
[0 1 4 5] [2 3]
[0 1 2 3] [4 5]

pandas との連携 (ダメなパターン)

KFold#split() で手に入るのは単なるインデックスなので、それを元に DataFrame から対象データを抽出しないといけない。このとき、次のようなコードを書いてしまうと一見すると上手くいっているように見えて後でハマることになる。

>>> from sklearn.model_selection import KFold
>>> kf = KFold(n_splits=3)
>>> for train, test in kf.split(df):
...     # これはやっちゃダメ！
...     train_df = df[df.index.isin(train)]
...     test_df = df[df.index.isin(test)]
...     # 結果確認 (一見すると上手くいっているように見える)
...     print('(train)', train_df)
...     print('(test)', test_df)
...     print('-----')
... 
(train)     name
2    Cat
3   Deer
4  Eagle
5    Fox
(test)      name
0     Ant
1  Beetle
-----
(train)      name
0     Ant
1  Beetle
4   Eagle
5     Fox
(test)    name
2   Cat
3  Deer
-----
(train)      name
0     Ant
1  Beetle
2     Cat
3    Deer
(test)     name
4  Eagle
5    Fox
-----

実は、上記はインデックスが 0 からの連番で振られているために動いているに過ぎない。その前提が変わると途端に動かなくなる。試しにインデックスの番号を変更してみよう。

>>> # インデックスを 0 から始まる連番から変える (以下なら 10, 20, 30...)
... df.index = df.index * 10 + 10

変更後のインデックスはこんな感じ。

>>> df
      name
10     Ant
20  Beetle
30     Cat
40    Deer
50   Eagle
60     Fox

先ほどと全く同じコードを使って動作を確認してみよう。

>>> kf = KFold(n_splits=3)
>>> for train, test in kf.split(df):
...     # OMG!!
...     train_df = df[df.index.isin(train)]
...     test_df = df[df.index.isin(test)]
...     # 結果確認
...     print('(train)', train_df)
...     print('(test)', test_df)
...     print('-----')
... 
(train) Empty DataFrame
Columns: [name]
Index: []
(test) Empty DataFrame
Columns: [name]
Index: []
-----
(train) Empty DataFrame
Columns: [name]
Index: []
(test) Empty DataFrame
Columns: [name]
Index: []
-----
(train) Empty DataFrame
Columns: [name]
Index: []
(test) Empty DataFrame
Columns: [name]
Index: []
-----

中身が全て空っぽになってしまっている！

pandas との連携 (正解)

上記のようなパターンでも動くようにするには DataFrame の絞り込みで DataFrame#iloc を使う。これなら、本来のインデックスの値ではなく中身の順序にもとづいたインデックスで絞り込みができる。

>>> kf = KFold(n_splits=3)
>>> for train, test in kf.split(df):
...     # これなら上手くいく
...     train_df = df.iloc[train]
...     test_df = df.iloc[test]
...     print('(train)', train_df)
...     print('(test)', test_df)
...     print('-----')
... 
(train)      name
30    Cat
40   Deer
50  Eagle
60    Fox
(test)       name
10     Ant
20  Beetle
-----
(train)       name
10     Ant
20  Beetle
50   Eagle
60     Fox
(test)     name
30   Cat
40  Deer
-----
(train)       name
10     Ant
20  Beetle
30     Cat
40    Deer
(test)      name
50  Eagle
60    Fox
-----