scikit-learnでデータを訓練用とテスト用に分割するtrain_test_split
scikit-learnのtrain_test_split()関数を使うと、NumPy配列ndarrayやリストなどを二分割できる。機械学習においてデータを訓練用（学習用）とテスト用に分割してホールドアウト検証を行う際に用いる。
ここでは以下の内容について説明する。
train_test_split()の基本的な使い方- 割合、個数を指定: 引数
test_size,train_size - シャッフルするかを指定: 引数
shuffle - 乱数シードを指定: 引数
random_state - 複数のデータを分割
- 層化抽出: 引数
stratify - 具体的な例（アイリスデータセット）
pandas.DataFrame,Seriesの場合
train_test_split()の基本的な使い方
train_test_split()にNumPy配列ndarrayを渡すと、二分割されたndarrayが要素として格納されたリストが返される。
import numpy as np
from sklearn.model_selection import train_test_split
a = np.arange(10)
print(a)
# [0 1 2 3 4 5 6 7 8 9]
print(train_test_split(a))
# [array([3, 9, 6, 1, 5, 0, 7]), array([2, 8, 4])]
print(type(train_test_split(a)))
# <class 'list'>
print(len(train_test_split(a)))
# 2
以下のように、アンパックでそれぞれ2つの変数に代入することが多い。
a_train, a_test = train_test_split(a)
print(a_train)
# [3 4 0 5 7 8 2]
print(a_test)
# [6 1 9]
例はnumpy.ndarrayだが、list（Python組み込みのリスト）やpandas.DataFrame, Series、疎行列scipy.sparseにも対応している。pandas.DataFrame, Seriesの例は最後に示す。
割合、個数を指定: 引数test_size, train_size
引数test_sizeでテスト用（返されるリストの2つめの要素）の割合または個数を指定できる。
デフォルトはtest_size=0.25で25%がテスト用、残りの75%が訓練用となる。小数点以下は切り上げとなり、上の例では10 * 0.25 = 2.5 -> 3となっている。
test_sizeには0.0 ~ 1.0の割合か、個数を指定する。
割合で指定。
a_train, a_test = train_test_split(a, test_size=0.6)
print(a_train)
# [9 1 2 6]
print(a_test)
# [5 7 4 3 0 8]
個数で指定。
a_train, a_test = train_test_split(a, test_size=6)
print(a_train)
# [4 2 1 0]
print(a_test)
# [7 6 3 9 8 5]
引数train_sizeで訓練用の割合・個数を指定することもできる。test_sizeと同様に、0.0 ~ 1.0の割合か、個数を指定する。
a_train, a_test = train_test_split(a, train_size=0.6)
print(a_train)
# [2 9 6 0 4 3]
print(a_test)
# [7 8 5 1]
a_train, a_test = train_test_split(a, train_size=6)
print(a_train)
# [9 3 0 8 7 1]
print(a_test)
# [5 6 4 2]
train_sizeは小数点以下切り捨て。
a_train, a_test = train_test_split(a, train_size=0.25)
print(a_train)
# [1 2]
print(a_test)
# [0 8 4 7 5 6 3 9]
これまでの例のように引数test_size, train_sizeのいずれかのみを指定した場合、他方の数はその残りになるが、それぞれを別途指定することも可能。
a_train, a_test = train_test_split(a, test_size=0.3, train_size=0.4)
print(a_train)
# [3 0 4 9]
print(a_test)
# [7 2 8]
a_train, a_test = train_test_split(a, test_size=3, train_size=4)
print(a_train)
# [9 7 0 4]
print(a_test)
# [3 8 5]
上の例のように訓練用とテスト用を合わせて100%以下でも問題ないが、100%を超える値を指定するとエラーとなる。
# a_train, a_test = train_test_split(a, test_size=0.8, train_size=0.7)
# ValueError: The sum of test_size and train_size = 1.500000, should be smaller than 1.0. Reduce test_size and/or train_size.
# a_train, a_test = train_test_split(a, test_size=8, train_size=7)
# ValueError: The sum of train_size and test_size = 15, should be smaller than the number of samples 10. Reduce test_size and/or train_size.
シャッフルするかを指定: 引数shuffle
これまでの例のように、デフォルトでは要素がシャッフルされて分割される。引数shuffle=Falseとするとシャッフルされずに先頭から順番に分割される。
a_train, a_test = train_test_split(a, shuffle=False)
print(a_train)
# [0 1 2 3 4 5 6]
print(a_test)
# [7 8 9]
乱数シードを指定: 引数random_state
シャッフルされる場合、デフォルトでは実行するたびにランダムに分割される。引数random_stateを指定して乱数シードを固定すると常に同じように分割される。
a_train, a_test = train_test_split(a, random_state=0)
print(a_train)
# [9 1 6 7 3 0 5]
print(a_test)
# [2 8 4]
機械学習のモデルの性能を比較するような場合、どのように分割されるかによって結果が異なってしまうため、乱数シードを固定して常に同じように分割されるようにする必要がある。
複数のデータを分割
これまでの例では1つのデータ（numpy.ndarray）を分割したが、train_test_split()は複数のデータを分割することができる。
以下の2つのデータを例とする。
X = np.arange(20).reshape(2, 10).T
print(X)
# [[ 0 10]
# [ 1 11]
# [ 2 12]
# [ 3 13]
# [ 4 14]
# [ 5 15]
# [ 6 16]
# [ 7 17]
# [ 8 18]
# [ 9 19]]
y = np.arange(10)
print(y)
# [0 1 2 3 4 5 6 7 8 9]
train_test_split()の引数に順に指定すると、以下のように分割される。2つのデータが対応して分割されているのが確認できる。
X_train, X_test, y_train, y_test = train_test_split(X, y)
print(X_train)
# [[ 7 17]
# [ 3 13]
# [ 0 10]
# [ 8 18]
# [ 6 16]
# [ 4 14]
# [ 2 12]]
print(X_test)
# [[ 5 15]
# [ 1 11]
# [ 9 19]]
print(y_train)
# [7 3 0 8 6 4 2]
print(y_test)
# [5 1 9]
3つ以上でも問題ない。
z = np.arange(10) * 10
print(z)
# [ 0 10 20 30 40 50 60 70 80 90]
X_train, X_test, y_train, y_test, z_train, z_test = train_test_split(X, y, z)
print(X_train)
# [[ 6 16]
# [ 9 19]
# [ 1 11]
# [ 2 12]
# [ 7 17]
# [ 0 10]
# [ 3 13]]
print(X_test)
# [[ 8 18]
# [ 4 14]
# [ 5 15]]
print(y_train)
# [6 9 1 2 7 0 3]
print(y_test)
# [8 4 5]
print(z_train)
# [60 90 10 20 70 0 30]
print(z_test)
# [80 40 50]
それぞれのデータのlen()またはshape[0]（最初の次元の大きさ）が一致していないとエラー。
y_mismatch = np.arange(8)
print(y_mismatch)
# [0 1 2 3 4 5 6 7]
# X_train, X_test, y_train, y_test = train_test_split(X, y_mismatch)
# ValueError: Found input variables with inconsistent numbers of samples: [10, 8]
層化抽出: 引数stratify
例えば教師あり学習では特徴行列（説明変数）と正解ラベル（目的変数）の2つのデータを用いる。
二値分類（2クラス分類）では正解ラベルは、例えば以下のように0, 1のいずれかになる。
y = np.array([0] * 5 + [1] * 5)
print(y)
# [0 0 0 0 0 1 1 1 1 1]
教師あり学習のためにデータを分割する場合、訓練用とテスト用の正解ラベルの比率は元のデータの正解ラベルの比率と一致していることが望ましいが、例えば以下のようにテスト用に0の要素が含まれていないといったことが起こり得る。
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=100)
print(y_train)
# [0 1 0 0 0 0 1 1]
print(y_test)
# [1 1]
引数stratifyに均等に分割させたいデータ（多くの場合は正解ラベル）を指定すると、そのデータの値の比率が一致するように分割される。
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=100,
stratify=y)
print(y_train)
# [1 1 0 0 0 1 1 0]
print(y_test)
# [1 0]
サンプル数が少ないとイメージしにくいので、次の具体例も参照されたい。
具体的な例（アイリスデータセット）
具体的な例として、アイリスデータセットを分割する。
150件のデータがSepal Length（がく片の長さ）、Sepal Width（がく片の幅）、Petal Length（花びらの長さ）、Petal Width（花びらの幅）の4つの特徴量を持っており、Setosa, Versicolor, Virginicaの3品種に分類されている。
load_iris()でデータを取得する。正解ラベルyには0, 1, 2の3種類が均等に含まれている。
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
data = load_iris()
X = data['data']
y = data['target']
print(X.shape)
# (150, 4)
print(X[:5])
# [[5.1 3.5 1.4 0.2]
# [4.9 3. 1.4 0.2]
# [4.7 3.2 1.3 0.2]
# [4.6 3.1 1.5 0.2]
# [5. 3.6 1.4 0.2]]
print(y.shape)
# (150,)
print(y)
# [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
# 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
# 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 2 2 2 2 2 2 2 2 2 2 2
# 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
# 2 2]
train_test_split()で以下のように分割できる。test_sizeやtrain_sizeは設定していないので、デフォルトの通り、訓練用75%とテスト用25%に分割される。サイズが大きいので形状shapeのみ示している。
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
print(X_train.shape)
# (112, 4)
print(X_test.shape)
# (38, 4)
print(y_train.shape)
# (112,)
print(y_test.shape)
# (38,)
テスト用の正解ラベルy_testを確認すると、各ラベルの数が不均等になっている。
print(y_test)
# [2 1 0 2 0 2 0 1 1 1 2 1 1 1 1 0 1 1 0 0 2 1 0 0 2 0 0 1 1 0 2 1 0 2 2 1 0
# 1]
print((y_test == 0).sum())
# 13
print((y_test == 1).sum())
# 16
print((y_test == 2).sum())
# 9
引数stratifyを指定すると、分割後の各ラベルの比率が元のデータの比率（この例では3種類が均等）と一致するように分割できる。以下の例のように分割後の個数によっては完全に一致しない（割り切れない）こともあるが、できる限り元の比率に近くなっていることが分かる。
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0, stratify=y)
print(y_test)
# [0 0 0 0 1 1 1 0 1 2 2 2 1 2 1 0 0 2 0 1 2 1 1 0 2 0 0 1 2 1 0 1 2 2 0 1 2
# 2]
print((y_test == 0).sum())
# 13
print((y_test == 1).sum())
# 13
print((y_test == 2).sum())
# 12
pandas.DataFrame, Seriesの場合
pandas.DataFrame, Seriesはそれぞれ二次元・一次元の配列と同様に分割できる。
ここでもアイリスデータセットを例とする。
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
data = load_iris()
X_df = pd.DataFrame(data['data'], columns=data['feature_names'])
y_s = pd.Series(data['target'])
print(X_df)
# sepal length (cm) sepal width (cm) petal length (cm) petal width (cm)
# 0 5.1 3.5 1.4 0.2
# 1 4.9 3.0 1.4 0.2
# 2 4.7 3.2 1.3 0.2
# 3 4.6 3.1 1.5 0.2
# 4 5.0 3.6 1.4 0.2
# .. ... ... ... ...
# 145 6.7 3.0 5.2 2.3
# 146 6.3 2.5 5.0 1.9
# 147 6.5 3.0 5.2 2.0
# 148 6.2 3.4 5.4 2.3
# 149 5.9 3.0 5.1 1.8
#
# [150 rows x 4 columns]
print(type(X_df))
# <class 'pandas.core.frame.DataFrame'>
print(X_df.shape)
# (150, 4)
print(y_s)
# 0 0
# 1 0
# 2 0
# 3 0
# 4 0
# ..
# 145 2
# 146 2
# 147 2
# 148 2
# 149 2
# Length: 150, dtype: int64
print(type(y_s))
# <class 'pandas.core.series.Series'>
print(y_s.shape)
# (150,)
引数などの使い方はnumpy.ndarrayの場合とまったく同じ。分割されたpandas.DataFrame, Seriesが返される。
X_train_df, X_test_df, y_train_s, y_test_s = train_test_split(
X_df, y_s, test_size=0.2, random_state=0, stratify=y_s
)
print(type(X_train_df))
# <class 'pandas.core.frame.DataFrame'>
print(X_train_df.shape)
# (120, 4)
print(type(X_test_df))
# <class 'pandas.core.frame.DataFrame'>
print(X_test_df.shape)
# (30, 4)
print(type(y_train_s))
# <class 'pandas.core.series.Series'>
print(y_train_s.shape)
# (120,)
print(type(y_test_s))
# <class 'pandas.core.series.Series'>
print(y_test_s.shape)
# (30,)
引数stratifyの効果を確認する。訓練用もテスト用も元のデータの比率に合わせて分割されていることが分かる。
print(y_train_s.value_counts())
# 2 40
# 1 40
# 0 40
# dtype: int64
print(y_test_s.value_counts())
# 2 10
# 1 10
# 0 10
# dtype: int64