Ax でベイズ最適化によるパラメータチューニングを可視化

7 min read読了の目安(約6900字

Ax は Facebook のオープンソース最適化ライブラリです。
カーネルや獲得関数を定義しなくても動くので、知識がなくてもベイズ最適化が使えます。
Ax にはインタラクティブのグラフを描く api があり、ベイズ最適化によるパラメータチューニングを可視化できます。

インストール

pip あるいは conda でインストールします。

pip install ax-platform jupyter

APIs

Ax の API には loop, service, developer の3つがあります。
中でも developer api はデータサイエンティスト、機械学習エンジニア、研究者向けに作られた api でカスタム性が高いのが特徴です。今回は developer api を使ってパラメータチューニングをします。

ベイズ最適化について

ベイズ最適化は目的関数が最大になるパラメータを効率よく探索する数理最適化手法です。
期待値が大きいところを採用する「活用」と、分散が大きいところを採用する「探索」を数値化した獲得関数が最大になるパラメータを推薦し、これを繰り返して逐次的に最適な条件を探します。
ベイズ最適化で膨大な時間がかかるパラメータチューニングをさっくりと済ませられます。

SVR のパラメータチューニング

ボストン住宅価格データセットを用いて SVR のパラメータチューニングをします。

import numpy as np
from ax import ParameterType, RangeParameter, SearchSpace, SimpleExperiment, modelbridge, models
from sklearn.svm import SVR
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split, KFold
from sklearn.datasets import load_boston
from ax.utils.notebook.plotting import render, init_notebook_plotting
init_notebook_plotting()

boston = load_boston()
X = boston.data
y = boston.target
X_train, X_test, y_train, y_test = train_test_split(X, y)
scaler = StandardScaler()
autoscale_X_train = scaler.fit_transform(X_train)
autoscale_X_test = scaler.transform(X_test)
kf = KFold()

次に SVR のハイパーパラメータ C, epsilon, gamma の探索領域を定義します。

svr_search_space = SearchSpace(
    parameters=[
        RangeParameter(name="C",
                       parameter_type=ParameterType.FLOAT,
                       lower=0.1, 
                       upper=100, 
                       log_scale=True
                      ), 
        RangeParameter(name="epsilon",
                       parameter_type=ParameterType.FLOAT, 
                       lower=0.001,
                       upper=100,
                       log_scale=True
                      ),
        RangeParameter(name="gamma",
                       parameter_type=ParameterType.FLOAT, 
                       lower=0.001,
                       upper=100,
                       log_scale=True
                      ),
    ]
)

ベイズ最適化でパラメータチューニングをします。
SVR のスコアを返す関数 svr_evaluation_function を事前に定義し、SimpleExperiment オブジェクトを生成します。

kfold_scores=[]
kfold_parameters=[]
for train_index, valid_index in kf.split(X_train):
    X2d_train = X_train[train_index]
    X2d_valid = X_train[valid_index]
    y2d_train = y_train[train_index]
    y2d_valid = y_train[valid_index]
    scaler_2d = StandardScaler()
    autoscale_X2d_train = scaler_2d.fit_transform(X2d_train)
    autoscale_X2d_valid = scaler_2d.transform(X2d_valid)
    # 目的関数を出力する関数
    def svr_evaluation_function(parameterization, weight=None):
        svr = SVR(C=parameterization["C"],
                  epsilon=parameterization["epsilon"],
                  gamma=parameterization["gamma"]
                 )
        svr.fit(autoscale_X2d_train, y2d_train)
        accuracy = svr.score(autoscale_X2d_valid, y2d_valid)
        return {"accuracy": (accuracy, 0.0)}
    exp = SimpleExperiment(name="test_svm",
                           search_space=svr_search_space,
                           evaluation_function=svr_evaluation_function,
                           objective_name="accuracy"
                          )
    # ベイズ最適化
    sobol = modelbridge.get_sobol(search_space=exp.search_space)
    print(f"Running Sobol initialization trials...")
    for _ in range(5):
        exp.new_trial(generator_run=sobol.gen(1))
    for i in range(15):
        print(f"Running GP+EI optimization tiral {i+1}/15...")
        gpei=modelbridge.get_GPEI(experiment=exp, data=exp.eval())
        exp.new_trial(generator_run=gpei.gen(1))
    best_objectives = exp.eval().df["mean"]
    bestParameter_score = np.max(best_objectives)
    bestParameter_trial = np.argmax(best_objectives)
    bestParameter = exp._trials[bestParameter_trial].arm.parameters
    kfold_scores.append(bestParameter_score)
    kfold_parameters.append(bestParameter)
# チュー二ングされたパラメータ
svr_parameter = kfold_parameters[np.argmax(kfold_scores)]
svr = SVR(**svr_parameter)
# テストデータでスコアを計算
svr.fit(autoscale_X_train, y_train)
svr_score = svr.score(autoscale_X_test, y_test)

sobol で最初の 5 点をランダムに決めて、20 回探索で最もスコアの高かったパラメータをベストパラメータとします。
exp.evalで過程が見れます。
best_objectives には探索した期待値が入っており、最も高いスコアをベストパラメータとしています。
gpei の名前はガウス過程回帰と獲得関数 EI から来ており, gpeiに詳細が記録されています。
グラフを描くときは、gpeiを参照することになります。

チューニングされたパラメータを用いてテストデータのスコアを計算します。

>>> svr_score
0.77991406620649

計算時間は 1 分 15 秒でした。

探索過程の可視化

jupyterlab で plotly をインラインで描けるように設定して ax でグラフを描きます。
詳細は以下のリンクを参照してください。
オフライン環境でplotlyを活用するための基礎(html出力、png出力、インライン描画)

先ほどの逐次最適化の過程をグラフにします。
横軸を探索回数、縦軸をベストスコアにしたグラフです。

from ax.plot.trace import optimization_trace_single_method

best_objectives = np.array([[trial.objective_mean*100 for trial in exp.trials.values()]])
best_objective_plot = optimization_trace_single_method(
y=np.maximum.accumulate(best_objectives, axis=1),
title="Model performance vs. # of iterations",
ylabel="Regression Accuracy, %",)
render(best_objective_plot)

17 回以降はスコアが変わらなかったので、チューニングができているはずです。

2 つのパラメータを軸に取って期待値と標準偏差を表したグラフです。

from ax.plot.contour import interact_contour

render(interact_contour(model=gpei, metric_name='accuracy'))

グラフを見ると期待値が大きな領域を重点的に探索しています。
ただ外れた領域も探索している傾向が見られます。
これは探索の過程で標準偏差が大きい領域が見つかったが、計算したところ期待値は小さく以降探索されなくなったのだと思います。

パラメータ 1 つ選択してスコアの期待値と標準偏差を表したグラフです。

from ax.plot.slice import plot_slice

render(plot_slice(gpei, "epsilon", "accuracy"))

まとめ

Ax で SVR のパラメータチューニングを行いました。sklearn のモデルにも簡単に組み込めて使いやすかったです。機械学習のパラメータチューニングにはもってこいです。

おまけ

グリッドサーチでパラメータチューニングをした結果です。

from sklearn.pipeline import make_pipeline
from sklearn.model_selection import GridSearchCV

pipe = make_pipeline(StandardScaler(), SVR())
C_grid = np.linspace(0.01, 100, 50)
gamma_grid = np.linspace(0.001, 100, 50)
epsilon_grid = np.linspace(0.001, 100, 50)
param_grid = [{pipe.steps[1][0] + "__C": C_grid, 
               pipe.steps[1][0] + "__gamma": gamma_grid, 
               pipe.steps[1][0] + "__epsilon": epsilon_grid
              }]
reg = GridSearchCV(pipe, param_grid=param_grid)
reg.fit(X_train, y_train)
score = reg.score(X_test, y_test)
>>> score
0.6963721758604782

計算時間は10分49秒でした。分割数が50点で少なめですが、ベイズ最適化よりもおよそ10倍時間がかかり、スコアは低くなっています。