scikit-learnでトレーニングデータとテストデータを分割する

機械学習させるときに用意したデータの9割を学習用データに、1割をテストデータに使ったりします。その分割を手軽にしてくれます。scikit-learnを使いたくてPython書いてるようなもんです。

from sklearn.cross_validation import train_test_split
import numpy as np

def main():

    _train_x, _train_y, _test_x, _test_y = train_test_split(
        np.array([
            [0, 1, 0, 1, 0, 1, 0, 1, 0, 1],
            [1, 1, 0, 1, 0, 1, 0, 1, 0, 1],
            [2, 1, 0, 1, 0, 1, 0, 1, 0, 1],
            [3, 1, 0, 1, 0, 1, 0, 1, 0, 1],
            [4, 1, 0, 1, 0, 1, 0, 1, 0, 1],
            [5, 1, 0, 1, 0, 1, 0, 1, 0, 1],
            [6, 1, 0, 1, 0, 1, 0, 1, 0, 1],
            [7, 1, 0, 1, 0, 1, 0, 1, 0, 1],
            [8, 1, 0, 1, 0, 1, 0, 1, 0, 1],
            [9, 1, 0, 1, 0, 1, 0, 1, 0, 1]
        ]),
        np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1]),
        test_size=0.1
    )
    print(_train_x)

if __name__ == "__main__":
    main()

test_size=0.1、で10%をテスト用データとして分割してくれます。

[[0 1 0 1 0 1 0 1 0 1]
 [8 1 0 1 0 1 0 1 0 1]
 [7 1 0 1 0 1 0 1 0 1]
 [4 1 0 1 0 1 0 1 0 1]
 [2 1 0 1 0 1 0 1 0 1]
 [5 1 0 1 0 1 0 1 0 1]
 [6 1 0 1 0 1 0 1 0 1]
 [1 1 0 1 0 1 0 1 0 1]
 [3 1 0 1 0 1 0 1 0 1]]

これだけでシャッフルして指定した割合で分割してくれるからありがたい。