【scikit-learn】データを訓練データとテストデータに分割するtrain_test_splitについて



1. 訓練データとテストデータの分割
機械学習のプロジェクトを進めるうえで、データを適切に分割することは非常に重要です。モデルを訓練するためのデータ(訓練データ)と、モデルの性能を評価するためのデータ(テストデータ)に分けることで、過学習を防ぎ、モデルの一般化性能を測ることができます。この記事では、Pythonのライブラリ「scikit-learn」を使用して、データを効率よく分割する方法について解説します。
1.1. データ分割の必要性
機械学習モデルは、過去のデータからパターンを学習し、新しいデータに適応することを目指します。しかし、モデルが訓練データにだけ最適化されてしまうと、新しいデータに対してうまく予測できなくなります。これを防ぐために、以下のようにデータを分割します。
- 訓練データ(Training Data) モデルを学習させるために使用するデータ。
- テストデータ(Test Data) 訓練されたモデルの性能を評価するために使用するデータ。
2. scikit-learnのtrain_test_split
scikit-learnでは、データ分割を簡単に行えるように、train_test_split
が用意されています。この関数を使うことで、データをランダムに訓練データとテストデータに分けることができます。

train_test_split
が便利なんでしょ?
test_size
で分割比率を指定できる。2.1. train_test_splitの基本構文
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
2.2. パラメータの説明
X
: 特徴量(説明変数)のデータy
: ターゲット(目的変数)のデータtest_size
: テストデータの割合(デフォルトは0.25)random_state
: データ分割時のランダムシード(結果を再現可能にするため)train_size
: 訓練データの割合(test_size
と合わせて指定可能)
3. 実際の利用例
3.1. データセットの準備
まずはデータセットを準備します。ここでは、scikit-learnで利用できるtoyデータセットである「Irisデータセット」を使用します。
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
# データセットの読み込み
iris = load_iris()
X = iris.data # 特徴量
y = iris.target # ラベル
3.2. 訓練データとテストデータに分割
# データの分割
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 分割結果の確認
print("訓練データのサイズ:", X_train.shape)
print("テストデータのサイズ:", X_test.shape)
3.3. 実行結果
ここでは、データ全体の20%がテストデータ、80%が訓練データに分割されました。
3.4. 分割比率を変更する
テストデータの割合を変更したい場合は、test_size
を調整します。
3.5. 例: テストデータを30%にする
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
ここでは、データ全体の30%がテストデータ、70%が訓練データに分割されました。
4. stratifyを使った分割(クラスの偏りを防ぐ)
分類問題では、データのクラス分布が偏らないようにすることが重要です。train_test_split
にはstratify
オプションがあり、これを指定すると分割後のクラス分布を元データに合わせることができます。


stratify
を使う。4.1. 例: クラス分布を気にしないで分割した場合
まず、stratify
オプションを指定せずにデータを分割してみます。以下のコードは、データを訓練データとテストデータに8:2の割合で分割し、その際のクラス分布を確認する例です。
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# クラス分布の確認
import numpy as np
print("訓練データのクラス分布:", np.bincount(y_train))
print("テストデータのクラス分布:", np.bincount(y_test))
このコードでは、NumPyのbincount
関数を使用してクラスごとの出現回数を数えています。bincount
は、整数の配列内で各値が何回現れるかを効率的に集計するための便利な関数です。
例えば、元のデータのクラス分布が不均一であった場合、このように分割すると訓練データとテストデータのクラス分布も偏りが生じる可能性があります。
4.2. 例: クラス分布を保った分割
次に、stratify
オプションを使用して、データ分割後も元のクラス分布が維持されるようにします。具体的には、分割時にstratify=y
と指定することで、ラベル(y
)の分布を保ちながら訓練データとテストデータを作成します。
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
# クラス分布の確認
import numpy as np
print("訓練データのクラス分布:", np.bincount(y_train))
print("テストデータのクラス分布:", np.bincount(y_test))
このコードでは、元データのラベル分布を基準として、訓練データとテストデータを適切に分割します。その結果、クラスの偏りがなくなり、データがバランス良く分けられます。