🔄

ML モデル実装に Strategy パターンを導入してみた

に公開

はじめに

Fusicのレオナです。
大学では機械学習・深層学習のモデル開発に取り組んでおり、入社後は新卒研修でソフトウェア開発を学んだあと、機械学習モデル開発案件に携わりました。

その中で強く感じたのが、精度改善そのものだけでなく、比較しやすい実装にしておくことの大切さです。

機械学習の開発では、PoCや初期検証の段階で複数のモデルを並行して試すことがよくあります。たとえば、RandomForest、LightGBM、ニューラルネットワークのように異なるモデルを切り替えながら、精度や学習コスト、実装のしやすさを見比べていきます。

最初のうちは if/elif で切り替える実装でも十分に見えます。ですが、モデルごとに初期化方法や学習処理が違い、検証パターンも増えてくると、切り替えのための分岐が徐々にコード全体の見通しを悪くしていきます。

この本ブログでは、モデルを切り替えながら検証するコードを整理する方法として、デザインパターンである、Strategyパターンを機械学習の実装に当てはめた例を書いてます。

Strategy パターンとは

Strategyパターンは、アルゴリズムや振る舞いを個別のクラスとして分離し、利用側から差し替え可能にするデザインパターンです。
GoF(Gang of Four)の「Design Patterns」で定義された振る舞いパターンの1つとして知られています。

機械学習の文脈に置き換えると、「どのモデルで学習・予測するか」を差し替え可能な戦略として扱うイメージです。

分岐でモデルを切り替えると何がつらいのか

まずは、Strategyパターンを使わずにモデルを切り替える例を見てみます。

def train_and_evaluate(model_type: str, X_train, y_train, X_test, y_test):
    if model_type == "random_forest":
        from sklearn.ensemble import RandomForestClassifier

        model = RandomForestClassifier(n_estimators=100, random_state=42)
        model.fit(X_train, y_train)
        predictions = model.predict(X_test)

    elif model_type == "lightgbm":
        import lightgbm as lgb

        model = lgb.LGBMClassifier(n_estimators=100, random_state=42, verbose=-1)
        model.fit(X_train, y_train)
        predictions = model.predict(X_test)

    elif model_type == "neural_network":
        import torch
        import torch.nn as nn

        # PyTorch のモデル定義・学習・推論をここにまとめて書く
        pass

    elif model_type == "svm":
        from sklearn.svm import SVC

        model = SVC(kernel="rbf")
        model.fit(X_train, y_train)
        predictions = model.predict(X_test)

    else:
        raise ValueError(f"Unknown model type: {model_type}")

    accuracy = (predictions == y_test).mean()
    return accuracy

この形は最初はわかりやすいのですが、モデルが増えるにつれて次のような問題が出てきます。

  • 固有の処理が1か所に集まりやすい

    • 学習器の初期化、学習、推論、評価までを1つの関数で持ち始めると、責務の境界が曖昧になります。
  • 比較実験を増やすたびに分岐が膨らむ

    • 新しいモデルや検証条件を追加するたびに既存コードへ手を入れることになり、差分の把握もしづらくなります。
  • 共通処理と個別処理の境目が見えにくい

    • 何が全モデル共通の流れで、何がモデル固有の実装なのかが読み取りづらくなります。

機械学習のコードでは、実験の回数が増えるほど「とりあえず動く実装」と「あとから手を入れやすい実装」の差が大きくなります。
そこで、モデルごとの差分を独立したクラスに切り出し、利用側は共通の操作だけを見るようにすると整理しやすくなります。

Strategy パターンでどう整理するか

今回の構成では、モデルの違いを「戦略」として切り出し、呼び出し側は共通インターフェースを通して扱います。

  1. Strategyインターフェースを定義する
    fit/predict の共通インターフェースを決める

  2. 各モデルをConcreteStrategyとして実装する
    RandomForest、LightGBM、PyTorch NN などを個別クラスに分ける

  3. Contextを用意する
    現在使う戦略を保持し、学習や予測を呼び出す

  4. 呼び出し側は戦略を差し替えるだけにする
    比較したいモデルを順番に渡して評価する

この形にしておくと、呼び出し側は「今どのモデルを使うか」だけに集中でき、各モデル固有の実装詳細は個別のクラスに閉じ込められます。

ディレクトリ構成

ml_strategy/
├── strategy.py
├── context.py
├── main.py
└── models/
    ├── __init__.py
    ├── random_forest.py
    ├── lightgbm_model.py
    ├── neural_network.py
    └── svm.py

Strategy インターフェースを定義する

今回は Python の Protocol を使って、fit/predict の共通インターフェースを定義します。

# strategy.py
from typing import Protocol
import numpy as np


class ModelStrategy(Protocol):
    """機械学習モデルの共通インターフェース"""

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        """モデルを学習する"""
        ...

    def predict(self, X: np.ndarray) -> np.ndarray:
        """予測結果を返す"""
        ...

Protocolを使うと、明示的に継承していなくても、fitpredictを備えたクラスであれば同じインターフェースとして扱えます。

今回は呼び出し側で fit の戻り値を使わないため、fitNoneを返す設計にしています。

各モデルを個別のクラスとして実装する

RandomForest

# models/random_forest.py
from sklearn.ensemble import RandomForestClassifier
import numpy as np


class RandomForestStrategy:
    """RandomForest を使う戦略"""

    def __init__(self, n_estimators: int = 100, random_state: int = 42):
        self.model = RandomForestClassifier(
            n_estimators=n_estimators,
            random_state=random_state,
        )

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        self.model.fit(X, y)

    def predict(self, X: np.ndarray) -> np.ndarray:
        return self.model.predict(X)

LightGBM

# models/lightgbm_model.py
import lightgbm as lgb
import numpy as np


class LightGBMStrategy:
    """LightGBM を使う戦略"""

    def __init__(self, n_estimators: int = 100, random_state: int = 42):
        self.model = lgb.LGBMClassifier(
            n_estimators=n_estimators,
            random_state=random_state,
            verbose=-1,
        )

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        self.model.fit(X, y)

    def predict(self, X: np.ndarray) -> np.ndarray:
        return self.model.predict(X)

PyTorch のニューラルネットワーク

PyTorch のように学習ループを自前で持つモデルでも、fit/predictの形にそろえておくと、利用側のコードを統一しやすくなります。

# models/neural_network.py
from typing import Optional

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset


class SimpleNN(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class NeuralNetworkStrategy:
    """PyTorch のニューラルネットワークを使う戦略"""

    def __init__(
        self,
        hidden_dim: int = 64,
        epochs: int = 50,
        lr: float = 0.001,
        batch_size: int = 32,
        random_state: int = 42,
    ):
        self.hidden_dim = hidden_dim
        self.epochs = epochs
        self.lr = lr
        self.batch_size = batch_size
        self.random_state = random_state

        self.model: Optional[SimpleNN] = None
        self.classes_: Optional[np.ndarray] = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def _set_seed(self) -> None:
        np.random.seed(self.random_state)
        torch.manual_seed(self.random_state)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(self.random_state)

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        self._set_seed()
        
        classes, y_encoded = np.unique(y, return_inverse=True)
        self.classes_ = classes

        input_dim = X.shape[1]
        output_dim = len(classes)

        self.model = SimpleNN(
            input_dim=input_dim,
            hidden_dim=self.hidden_dim,
            output_dim=output_dim,
        ).to(self.device)

        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
        criterion = nn.CrossEntropyLoss()

        dataset = TensorDataset(
            torch.tensor(X, dtype=torch.float32),
            torch.tensor(y_encoded, dtype=torch.long),
        )
        loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)

        self.model.train()
        for _ in range(self.epochs):
            for X_batch, y_batch in loader:
                X_batch = X_batch.to(self.device)
                y_batch = y_batch.to(self.device)

                optimizer.zero_grad()
                logits = self.model(X_batch)
                loss = criterion(logits, y_batch)
                loss.backward()
                optimizer.step()

    def predict(self, X: np.ndarray) -> np.ndarray:
        if self.model is None or self.classes_ is None:
            raise RuntimeError("fit() を先に呼び出してください")

        self.model.eval()
        X_tensor = torch.tensor(X, dtype=torch.float32).to(self.device)

        with torch.no_grad():
            logits = self.model(X_tensor)
            predicted_indices = torch.argmax(logits, dim=1).cpu().numpy()

        return self.classes_[predicted_indices]

PyTorch 側だけ学習の中身が複雑でも、利用側には同じ fit/predict として見せられることです。

Context を用意する

Context は、現在どの戦略を使うかを保持し、共通のワークフローを提供する役割を持ちます。

# context.py
import numpy as np
from strategy import ModelStrategy


class MLContext:
    """現在の戦略を使って学習・予測・評価を行う"""

    def __init__(self, strategy: ModelStrategy):
        self._strategy = strategy

    @property
    def strategy(self) -> ModelStrategy:
        return self._strategy

    @strategy.setter
    def strategy(self, strategy: ModelStrategy) -> None:
        """実行時に戦略を差し替える"""
        self._strategy = strategy

    def train(self, X: np.ndarray, y: np.ndarray) -> None:
        self._strategy.fit(X, y)

    def predict(self, X: np.ndarray) -> np.ndarray:
        return self._strategy.predict(X)

    def evaluate(self, X: np.ndarray, y: np.ndarray) -> float:
        predictions = self.predict(X)
        return float((predictions == y).mean())

利用側のコード

Iris データセットを使って、複数モデルを順番に学習・評価してみます。

# main.py
from typing import Dict

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

from context import MLContext
from models.lightgbm_model import LightGBMStrategy
from models.neural_network import NeuralNetworkStrategy
from models.random_forest import RandomForestStrategy


def main() -> None:
    # データの準備
    X, y = load_iris(return_X_y=True)
    X_train, X_test, y_train, y_test = train_test_split(
        X,
        y,
        test_size=0.3,
        random_state=42,
        stratify=y,
    )

    # 比較したい戦略を定義
    strategies = {
        "RandomForest": RandomForestStrategy(n_estimators=100),
        "LightGBM": LightGBMStrategy(n_estimators=100),
        "NeuralNetwork": NeuralNetworkStrategy(hidden_dim=64, epochs=100),
    }

    context = MLContext(strategy=RandomForestStrategy())
    results: Dict[str, float] = {}

    for name, strategy in strategies.items():
        context.strategy = strategy
        context.train(X_train, y_train)
        accuracy = context.evaluate(X_test, y_test)
        results[name] = accuracy
        print(f"{name}: {accuracy:.4f}")

    best_model = max(results, key=results.get)
    print(f"\nBest model: {best_model} ({results[best_model]:.4f})")


if __name__ == "__main__":
    main()

出力形式は、たとえば次のようになります。

RandomForest: 0.9333
LightGBM: 0.9556
NeuralNetwork: 0.9778

Best model: NeuralNetwork (0.9778)

これはあくまで実行結果の一例です。
seed や実行環境によって値は変わることがあります。

新しいモデルを追加する

Strategy パターンの良いところは、既存のContextや他の戦略を変更せずに拡張しやすいことです。

たとえば SVM を追加したい場合は、新しい戦略クラスを1つ増やします。

# models/svm.py
from sklearn.svm import SVC
import numpy as np


class SVMStrategy:
    """SVM を使う戦略"""

    def __init__(self, kernel: str = "rbf"):
        self.model = SVC(kernel=kernel)

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        self.model.fit(X, y)

    def predict(self, X: np.ndarray) -> np.ndarray:
        return self.model.predict(X)

利用側では辞書に1行追加するだけです。

from models.svm import SVMStrategy

strategies["SVM"] = SVMStrategy(kernel="rbf")

このとき、context.py や既存の RandomForestStrategy/LightGBMStrategy/ NeuralNetworkStrategyを修正する必要はありません。

最後に

今回は、機械学習のモデル比較コードに Strategy パターンを当てはめる例を紹介しました。モデル切り替えの責務を分離するだけでも、実験コードの見通しや拡張のしやすさはかなり変わります。

大学生時代にモデルの実装をもう少し考えて実装すればよかったと後悔してます。

Fusic 技術ブログ

Discussion