sklearn の train_test_split でデータの順番の情報を保持しておく #Python

sklearn の train_test_split でデータの順番の情報を保持しておく

この記事ではscikit-learn0.18.1を利用しています。

はじめに

データを学習用とテスト用に分割する train_test_split は非常に便利なのですが、
時系列データなど「順番」に意味がある場合、分割後もその情報を持っていたいケースがあるかと思います。

シンプルなデータであれば、順番の情報を持った上で分割し、後から取り除く・・・などあるのですが、
データが画像のような複雑な場合はうまくいかない。どうしよう・・・と思っていたら、
以下のようなシンプルな方法でできました。

データの分割

# coding:utf-8

import numpy as np
from sklearn.model_selection import train_test_split

X = np.arange(10).reshape((5, 2))
Y = np.array([1,1,0,1,0])
indices = np.array(range(X.shape[0]))

X_train, X_test, Y_train, Y_test, indices_train, indices_test = train_test_split(X, Y, indices, test_size=0.33, random_state=111)

indices で順番の情報を持っておいて、これもまとめて分割しています。

keras で画像データを解析していて、この問題に直面しました。
データフレームにすればインデックス番号を利用できると思ったけど、多次元データのデータフレーム化は
なかなかうまくいかず、numpy でいろいろやろうと思ったけどうまくいかず。
で、調べていたら上のようなシンプルなやり方がありました。