Pythonのscikit-learnでデータセットを訓練データとテストデータに分割する方法を解説

scikit-learn(sklearnとも呼ばれます)は、Pythonで機械学習アルゴリズムを実装するために広く利用されているライブラリです。統計モデリングのための多彩なツールを備えており、強力かつ堅牢であることが特徴です。

分類、回帰、クラスタリング、次元削減など、さまざまなタスクをPythonの安定したインターフェースを通じて実行できます。また、NumPy・SciPy・Matplotlibの上に構築されている点も大きな特長です。

なぜデータを分割する必要があるのか

入力データを機械学習アルゴリズムに渡す前に、訓練データセットとテストデータセットに分割しておく必要があります。

選択したモデルにデータを適合させると、モデルは訓練データから学習を行います。この過程で、モデルはデータのパターンを把握し、未知の新しいデータにも対応できるよう汎化することを学びます。テストデータセットは、モデルの訓練中には一切使用しません。

すべてのハイパーパラメータの調整が完了し、最適な重みが設定された後、初めてテストデータセットを機械学習アルゴリズムに与えます。このテストデータは、アルゴリズムが新しいデータに対してどの程度うまく汎化できるかを検証するために使われます。

それでは、scikit-learnライブラリを使ってデータを分割する具体的な方法を見ていきましょう。

コード例

from sklearn.datasets import load_iris
my_data = load_iris()
X = my_data.data
y = my_data.target
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=2
)
print("訓練データの特徴量の次元")
print(X_train.shape)
print("テストデータの特徴量の次元")
print(X_test.shape)
print("訓練データのターゲット値の次元")
print(y_train.shape)
print("テストデータのターゲット値の次元")
print(y_test.shape)

出力結果

訓練データの特徴量の次元
(120, 4)
テストデータの特徴量の次元
(30, 4)
訓練データのターゲット値の次元
(120,)
テストデータのターゲット値の次元
(30,)

コードの解説

• まず、必要なパッケージをインポートします。
• この例では、アヤメ(Iris)データセットを環境に読み込みます。
• データセットから特徴量(X)とターゲット値(y)を分離します。
• train_test_split関数を使用し、訓練データとテストデータを80%対20%の比率で分割します。
• つまり、全データの20%は、モデルが新しいデータにどれだけうまく汎化できるかを確認するために確保されます。
• test_size=0.2でテストデータの割合を指定し、random_state=2を設定することで、何度実行しても同じ分割結果が得られるようになり、再現性が保証されます。
• 最後に、分割後の各データセットの次元をコンソールに出力して確認します。150件のサンプルのうち、120件が訓練用、30件がテスト用に割り当てられていることがわかります。

補足:層化抽出(stratify)の活用

分類問題では、train_test_splitの引数に「stratify=y」を指定すると、クラスごとのサンプル比率が訓練データとテストデータの両方で維持されます。特にクラス間のデータ数に偏りがある場合、このオプションを使うことで、より信頼性の高いモデル評価が可能になります。