以前、Optunaによる最適化結果を手軽に確認できるWebダッシュボードを開発・公開しました。公開からすでに半年以上が経過し、現在は公式に利用を推奨されるようになりました。Google Summer of Codeなどを通してcontributorやcommitterも増えつつある一方で、設計や実装に関しては資料を残してきませんでした（※1）。本記事ではダッシュボードの紹介をするとともに、開発に興味がある方向けに開発に役立つ情報をまとめておきます。

GitHub: https://github.com/optuna/optuna-dashboard

optuna-dashboardとは？

optuna-dashboardはOptunaによるハイパーパラメータの最適化結果をWebブラウザ上で簡単に確認できるツールです（※2）。試行結果だけではなく、ハイパーパラメータの重要度や相関関係なども確認できます。

インストールや使い方はとても簡単です。Optunaの試行結果をRDB(SQLite3など)に保存し、そのStorage URLを指定するだけです。OptunaのRDBバックエンドを利用したことのある方であれば迷わずに使えるでしょう。

$ pip install optuna-dashboard
$ optuna-dashboard sqlite:///db.sqlite3

実行すると次のように、各Trialでサンプルされたハイパーパラメーターや評価値が手軽に確認できます。ハイパーパラメータ間の相関関係、重要度もあわせて表示されるため、探索空間の見直し等にも有用です。

また公式に推奨している方法ではありませんが、OptunaのIn-Memory Storageを使っている方であっても、次のように呼び出すことも可能です（APIは変更する可能性もあるのでご注意ください）。

import optuna
from optuna_dashboard.app import create_app


def objective(trial):
    x1 = trial.suggest_float("x1", 0, 10)
    x2 = trial.suggest_float("x2", 0, 10)
    return (x1 - 2) ** 2 + (x2 - 5) ** 2


if __name__ == "__main__":
    storage = optuna.storages.InMemoryStorage()
    study = optuna.create_study(
        study_name="single-objective",
        storage=storage
    )
    study.optimize(objective, n_trials=100)

    app = create_app(storage)
    app.run(...)

optuna-dashboardの設計

optuna-dashboardは、TypeScript + React製のSPA(Single Page Application)とBottle製のJSON APIサーバーの2つからなります。

設計はとてもシンプルです。最適化の試行結果はAPIポーリングにより定期的に同期されます。デフォルトでは最大で10秒程度の同期遅延が発生しますが、ほとんどのユースケースにおいては問題にならないでしょう（※3）。またJSON APIのエンドポイントも上の図に示す5つだけです。

optuna-dashboardのWebフロントエンドのコードは、もともとGoptunaのために実装しました（※4）。SPAとして実装されているのはこれが理由です。移植性を踏まえるとこの選択は正解でしたが、OptunaのユーザーやPython開発者はかならずしもWebフロントエンド技術に詳しいわけではないため、開発に興味があってもここの理解が大変かもしれません。以降はWebフロントエンドまわりを中心に知っておくと開発役立つことを解説します。

React Routerによるクライアントサイドルーティング

SPAとして実装したため、ルーティングは基本的にクライアントサイドで行われます。クライアントサイドのルーティングにはReact Routerを使用しています。クライアントサイドでのルーティングは、BrowserRouterを使う方法とHashRouterを使う方法があります。BrowserRouterが /dashboard/studies/1 のようにサーバーサイドでのルーティングと同じようなURL設計で行うのに対し、HashRouterは  /#dashboard/studies/1 のようにURLのハッシュ部分の文字列を使用します。

今回はBrowserRouterを利用しています。BrowserRouterを使う際には、サーバー側のルーティング機能とコンフリクトしないように気をつける必要があります。今回は/dashboard のURLパスから始まるリクエストを受け取った際には、次のようにSPAのJSファイルを読み込むHTMLファイルを返すことで、それ以降のルーティングをすべてクライアントサイドでルーティングにまかせています（コードはこちら）。

# Accept any following paths for client-side routing
@app.get("/dashboard<:re:(/.*)?>")
def dashboard() -> BottleViewReturn:
    response.content_type = "text/html"
    return INDEX_HTML

Recoilによる状態管理

Reactのステート管理には、Recoilを利用しています。APIから取得したStudyやTrialの一覧はさまざまなページから参照されます。ページ遷移の前後でAPIを叩き直すことを防ぐためのナイーブな解決策として、React Routerよりも上位のコンポーネントのstateに保存し、propsとして引き渡していく方法があります。しかし上位のコンポーネントのstateに多くの情報を詰め込むと、stateの一部が更新されたときにそれ以下の多くのコンポーネントが再レンダリングされてしまい、Reactの良さを活かせません。

この問題を防ぐ方法としてreact-reduxなどのライブラリーを使うことも可能ですが、やや重厚で比較的学習コストが大きいように感じています。RecoilはまだExperimentalなソフトウェアで、APIも安定していないようですが、Reactユーザーであれば手軽に扱えるライブラリーです。今後もAPIが変わる可能性があるため使い方をここでは解説しません。詳細は公式ドキュメントを読んでみてください。

Plotly.jsによる可視化

Optunaの可視化機能は、PlotlyのPythonライブラリーを使って optuna.visualizationモジュール に実装されています。optuna-dashboardは、この処理をPlotlyのJavaScriptライブラリーを使って再実装しています。この作業はやや大変で、完璧に移植できていない可視化機能もあります。もし表示崩れや問題になるケースを見つけたら、ぜひPRをくださると嬉しいです。

Plotly.jsを利用するにあたって、Reactコンポーネントとして扱うための薄いラッパーライブラリーもあるようです。しかしそちらを使用するとplotly.js-distが使えず依存関係が大幅に増えることや、使わなくともそれほど大変ではないことを踏まえ利用していません。Reactコンポーネント側では <div id=”graph-edf”> のようなdivタグを表示し、Plotly.jsを直接呼び出しています（こちらも結局やっていることはPlotlyのReactライブラリーも同じようです）。

React Testing LibraryとJestによるユニットテスト

optuna-dashboardのユニットテストは、Google Summer of Codeで参加したHuzi Cheng(@chenghuzi)さんが追加してくれました。テストランナーおよびAssertionライブラリーとしてJestを使用し、ReactコンポーネントのレンダリングやDOMへのアクセスにはreact-testing-libraryを使用しています。

it("Filter rows of DataGrid", () => {
  interface DummyAttribute {
    id: number
    key: string
    value: number
  }
  const dummyAttributes = [
    { id: 1, key: "foo", value: 1 },
    { id: 2, key: "bar", value: 1 },
    { id: 3, key: "bar", value: 2 },
    { id: 4, key: "foo", value: 2 },
    { id: 5, key: "foo", value: 3 },
  ]
  const columns: DataGridColumn[] = [
    { field: "key", label: "Key", filterable: true },
    {
      field: "value",
      label: "Value",
      sortable: true,
    },
  ]

  const { queryAllByText } = render(
    <DataGrid
      columns={columns}
      rows={dummyAttributes}
      keyField={"id"}
    />
  )
  expect(queryAllByText("bar").length).toBe(2)

  // Filter rows by "foo"
  fireEvent.click(queryAllByText("foo")[0])
  expect(queryAllByText("foo").length).toBe(3)
  expect(queryAllByText("bar").length).toBe(0)
})

みなさんが開発をする中でユニットテストの結果がなぜかFAILしているというときには、 screen.logTestingPlaygroundURL() を呼ぶとブラウザ上でDOMの状態をチェックできるのでぜひ開発の際は活用してください。

Pyppeteer(with Headless Chrome)によるVisual Regression Testing

Plotly.jsによるグラフ描画が主な処理になるため、ユニットテストで保証できる振る舞いはそれほど多くありません。グラフの描画に必要な情報もやや複雑で、OptunaのFrozenTrialという試行情報を各テストケースで用意するのは大変です。また可視化用のコンポーネントは多く、そのすべてにおいて網羅的なユニットテストを確保することは現実的ではありませんでした。

そこでPyppeteerとHeadless Chromeを使った簡易的なVisual Regression Testingを用意しています。実際にさまざまな目的関数をOptunaで評価して、その評価結果の入ったInMemoryStorageをもとにDashboardを起動、Pyppeteerでアクセスします (全体のコードは こちら)。実行すると↓の画像のようにスクリーンショットが生成されるため、様々な目的関数に対する表示結果をひと目で確認できるようになりました。

この方法では​​、実際にOptunaで最適化を回し、そのStorage情報を使ってJSON APIサーバーを立ち上げているので、FrozenTrialのfixtureを手動で頑張って用意する必要はありません。Optunaは最適化の過程で探索空間が変化することがあります。さまざまな探索空間が存在しうる状態でダッシュボードが壊れないことを確認するのは大変な作業です。過去には探索空間が変化するStudyを表示する際にクラッシュしてしまうバグを見逃し、PRをマージしてしまったこともありましたが、このテストにより確認が容易になりました。

但しPyppeteerは開発がほぼ止まりつつあるソフトウェアなので、あまりヘビーに使うのは避けたほうがいいかもしれません。今後動かなくなるタイミングがくればSelenium等に移行する予定です。

おわりに

本記事ではoptuna-dashboardの設計や仕様技術について解説しました。Optunaの利用ユーザーの多くはWebフロントエンド技術に親しみがないかもしれませんが、本記事をきっかけに開発に参加してくれる方が増えてくれると嬉しいです。

脚注

※1: プログラムのビルドやテスト・Linterの実行方法については DEVELOPMENT.md に残しています。

※2: Optunaは以前からBokeh CLIを使った簡易的なWebダッシュボードを提供していました。こちらは実装がシンプルだった一方で機能の拡張が難しく、現在はすでにDeprecatedされています。

※3: 過去にはWebSocketを利用することや、サーバー側でグラフを画像としてレンダリングしてフロントエンド側で表示することを提案されたこともありましたが、設計をシンプルに保ちつつ、ダッシュボードのインタラクティブ性を損なわないためにはこの設計が最もバランスがいいと思います。

※4: Goptunaには多目的最適化やハイパーパラメータの重要度などいくつかの機能がまだサポートされていません。そのためoptuna-dashboardのほうがGoptunaのWebダッシュボードよりも開発が進んでいます。

ハイパーパラメーター最適化は、機械学習モデルがその性能を発揮するために重要なプロセスの1つです。Pythonのハイパーパラメーター最適化ライブラリとして有名な Optuna [1] は、様々な最適化アルゴリズムに対応しつつも、使いやすく設計的にも優れたソフトウェアです。本記事ではOptunaの内部実装についてソフトウェア的な側面を中心に解説します。

Optunaの内部実装を理解するためには、主要コンポーネントの役割と全体の動作の流れを押さえる必要があります。しかしOptunaの開発は活発で、コード量も多くなり、全体の流れをコードから読み取ることは難しくなってきました。そこで今回Minitunaという小さなプログラムを用意しました。Minitunaには全部で3つのversionがあり、それぞれ100行、200行、300行とコード量が増えていきます。最終的にできあがるのは300行程度の小さなプログラムですが、実用的な枝刈りアルゴリズムを備えていて結構本格的です。

Minituna: Simplified Optuna implementation for new contributors.
https://github.com/CyberAgentAILab/minituna

Minitunaの各バージョンはそれぞれ次のステップでOptunaの設計を理解してもらうことを意図して実装しています。

1. Optunaの主要コンポーネントとそれらの呼び出しの流れを理解する
2. categoricalやinteger、loguniformの扱い方を理解する
3. 枝刈りのAPIとMedian stopping ruleのアルゴリズムを理解する

Minitunaを写経して設計を理解できればOptunaのコードを読み進めることはそれほど難しくないでしょう。本記事ではMinitunaのコードから簡単に読み取れることはあえて扱いませんが、その代わりMinitunaのコードリーディングや写経を進める際に手がかりになるヒントやOptunaとの違いなどを補足して解説していきます。

minituna_v1: Trial, Study, Sampler, Storageの役割と流れ

Minituna v1は100行程度のプログラムです。非常に小さなプログラムですが、主要なコンポーネントはすでに実装されており、次のようなプログラムが動作します。ちなみに今回用意したexampleはすべてOptunaとの互換性があり、import文を import optuna as minituna のように切り替えても問題なく動作します。

import minituna_v1 as minituna

def objective(trial: minituna.Trial) -> float:
    x = trial.suggest_uniform("x", 0, 10)
    y = trial.suggest_uniform("y", 0, 10)
    return (x - 3) ** 2 + (y - 5) ** 2

if __name__ == "__main__":
    study = minituna.create_study()
    study.optimize(objective, 10)
    print("Best trial:", study.best_trial)

Optunaではこのように目的関数を定義します。目的関数はTrialオブジェクトを受け取り、評価値を返す関数です。この例では目的関数を10回呼び出し、その中で得られた最も良い評価値(ここでは最小化問題を考えます)とそのパラメーターを出力します。

それではminituna_v1のコードリーディングを進めてみてください。minituna_v1には5つのクラスが定義されています。これらはOptunaのコードにも頻繁に登場する重要なコンポーネントです。これらの役割と呼び出しの流れをコードから確認してみましょう。

• Study: 1つの最適化タスクに関する一連の情報を管理するコンポーネント
  ある最適化タスクに対して、どのアルゴリズムを使うのか(Sampler)、どこに試行結果を保存するのか(Storage)といった情報をすべてこのStudyコンポーネントで管理します。 create_study() 関数の引数で指定します。
• Trial: 各試行(Trial)に対応するコンポーネント
  目的関数はこのTrialオブジェクトが提供するAPIを経由してOptunaからパラメーターをサンプルしたり、中間評価値を報告して枝刈りを行います。
• Storage: 最適化の試行結果を保存するコンポーネント
  このコンポーネントのおかげで、OptunaはRDB storageなどにも対応でき、試行結果の永続化や分散最適化が可能になっています。
• FrozenTrial: Storageレイヤー上での各試行(Trial)の表現
  各試行(Trial)において目的関数の評価に使ったパラメーターや評価値を保持するコンポーネントです。たとえばRDB storageなどを利用する際も、SQLでDBから取り出した情報をこのFrozenTrialオブジェクトに詰め込みます。
• Sampler: 次に評価するパラメーターを選択するアルゴリズムを実装するコンポーネント
  どのパラメーターを評価すればより良い評価値を得られるかを求めるアルゴリズムを実装するコンポーネントです。ベイズ最適化や進化戦略アルゴリズムにおけるハイパーパラメータのサンプリングはこのSamplerコンポーネントに定義します。Minitunaでは説明を簡単にするためRandom samplingのみ実装しています。Optunaがサポートしているベイズ最適化や進化戦略のアルゴリズムの詳細については、本記事で詳しく扱いません。弊社AI Lab研究員 野村のPyCon JP 2019の発表 機械学習におけるハイパーパラメータ最適化の理論と実践をご覧ください。

minituna_v2: CategoricalやInteger、LogUniformの扱い方

minituna_v2では suggest_uniform() (一様分布からの実数パラメーターのsample)に加えて次のSuggest APIをサポートしています。

• suggest_categorical(name, choices=[…]) : 質的パラメーターをサンプルする
• suggest_int(name, low, high) : 整数パラメーターをサンプルする
• suggest_loguniform(name, low, high) : 対数スケールの空間から実数パラメーターをサンプルする

このおかげで、次のような目的関数も最適化できます。

def objective(trial):
    iris = sklearn.datasets.load_iris()
    x, y = iris.data, iris.target

    classifier_name = trial.suggest_categorical("classifier", ["SVC", "RandomForest"])
    if classifier_name == "SVC":
        svc_c = trial.suggest_loguniform("svc_c", 1e-10, 1e10)
        classifier_obj = sklearn.svm.SVC(C=svc_c, gamma="auto")
    else:
        rf_max_depth = trial.suggest_int("rf_max_depth", 2, 32)
        classifier_obj = sklearn.ensemble.RandomForestClassifier(
            max_depth=rf_max_depth, n_estimators=10
        )

    score = sklearn.model_selection.cross_val_score(
        classifier_obj, x, y, n_jobs=-1, cv=3
    )
    accuracy = score.mean()
    return 1 - accuracy # 最小化問題に変えています

minituna_v2 のコードを理解するうえで押さえておくべきポイントは、すべてのパラメーターがStorage上ではfloatで表現される点です。先程の例において、質的パラメーターは “SVC” や “RandomForest” のように文字列となりますが、storage上ではこれらもfloatで表現されます。そのために次の抽象基底クラスを導入します。

class BaseDistribution(abc.ABC):
    @abc.abstractmethod
    def to_internal_repr(self, external_repr: Any) -> float:
        ...

    @abc.abstractmethod
    def to_external_repr(self, internal_repr: float) -> Any:
        ...

各パラメーターはinternal_reprとexternal_reprという2つの表現があります。internal_reprはストレージ上でのパラメーター表現でありfloat値です。external_reprは実際に目的関数で扱う表現なので文字列だったり整数値だったりします。実際に動かしてみるのが理解に繋がるでしょう。

>>> import minituna_v2 as minituna
>>> distribution = minituna.CategoricalDistribution(choices=["SVC", "RandomForest"])
>>> distribution.to_internal_repr("SVC")
0
>>> distribution.to_internal_repr("RandomForest")
1
>>> distribution.to_external_repr(0)
'SVC'
>>> distribution.to_external_repr(1)
'RandomForest'

internal_repr と external_repr の変換にはdistributionオブジェクトが必要です。そのため、distributionオブジェクトはstorageにも保存します。FrozenTrialにdistributionsフィールドが追加されているのはそれが理由です。

MinitunaとOptunaはSuggest APIに関していくつか異なる点があります。

1.  OptunaにはDiscreteUniformDistributionやIntLogUniformDistributionも存在します。DiscreteUniformDistributionはIntUniformDistributionと離散化の幅が異なるだけで本質的に同じです。IntLogUniformDistributionはLogUniformDistributionと同じような対応をIntUniformDistributionに加えただけです。そのため、minituna_v2を写経すれば簡単に読みすすめることができます。
2. Optunaには、 suggest_float(name, low, high, step=None, log=False) というAPIが用意されています。これは suggest_uniform(low, high) や suggest_loguniform(low, high) のuniformという単語が少し分かりにくいことや、suggest_intとの一貫性を考慮して追加されました。suggest_floatは無事に導入されたので、suggest_uniformやsuggest_loguniform、suggest_discrerte_uniformはDeprecateする選択肢もあると思いますが、その辺りの議論は止まっています。
3. DistributionをRDBStorageやRedisStorageに保存する際には何らかの形式にserializeする必要があり、OptunaではJSONにserializeして格納しています。つまりCategoricalパラメーターとして選択できる型は、 CategoricalChoiceType = Union[None, bool, int, float, str] のようにJSON serialize可能なオブジェクトに限定されることを意味しています。

Storageレイヤーをしっかり理解する

minituna_v2まで読み進めた方は、Optunaのコードリーディングに移るのもいいでしょう。minituna_v3の解説に入る前に、Storageレイヤーの補足解説をはさんでおきます。

自分がOptunaのコードリーディングをはじめたとき、最初にしっかり理解を深めたのはStorageレイヤーでした。どういう情報がストレージに保存されているかを整理すれば、各機能を実装するためにOptunaがどういうことをしないといけないのか自然と想像がついてくるからです。MinitunaとOptunaの設計の違いも、Storageレイヤーのコードを読めばだいたい把握できるでしょう。Storargeレイヤーの理解には、RDBStorageのSQLAlchemy モデル定義を確認するのがおすすめです。

class StudyModel(BaseModel):
    __tablename__ = "studies"
    study_id = Column(Integer, primary_key=True)
    study_name = Column(String(MAX_INDEXED_STRING_LENGTH), ...)
    direction = Column(Enum(StudyDirection), nullable=False)


class TrialModel(BaseModel):
    __tablename__ = "trials"
    trial_id = Column(Integer, primary_key=True)
    number = Column(Integer)
    study_id = Column(Integer, ForeignKey("studies.study_id"))
    state = Column(Enum(TrialState), nullable=False)
    value = Column(Float)
    datetime_start = Column(DateTime, default=datetime.now)
    datetime_complete = Column(DateTime)


class TrialParamModel(BaseModel):
    __tablename__ = "trial_params"
    __table_args__ = (UniqueConstraint("trial_id", "param_name"),)
    param_id = Column(Integer, primary_key=True)
    trial_id = Column(Integer, ForeignKey("trials.trial_id"))
    param_name = Column(String(MAX_INDEXED_STRING_LENGTH))
    param_value = Column(Float)
    distribution_json = Column(String(MAX_STRING_LENGTH))

Storageレイヤーのコードを読む際の、補足事項は次のとおりです。

• MinitunaのStorageはコードをシンプルにするため、1つのStudyに関する情報しか保持できませんが、実際のOptunaのコードは複数のStudyに対応しています。
• 複数のStudyに対応しているため、TrialModelのtrial_idフィールドは必ずしもStudy内で単調に増加するフィールドではありません。アルゴリズムを実装するうえで、Study内で 1, 2, 3, … と単調に増加する識別子があると便利なので、Optunaには number フィールドが導入されました。
• TrialStateには、 PRUNED や WAITING という状態があります。前者はminituna_v3で解説する枝刈り機能、後者は enqueue_trial() の機能を実装するために追加されました。

minituna_v3: Median stopping ruleによる枝刈り

minituna_v3は300行程度のプログラムです。枝刈り(早期停止)アルゴリズムをサポートしています。枝刈りは必ずしもOptunaの利用者全員が使う機能ではないので、興味のある方だけ取り組むとよいでしょう。枝刈り機能はつぎのように利用します。

def objective(trial):
    clf = MLPClassifier(
        hidden_layer_sizes=tuple(
            [trial.suggest_int("n_units_l{}".format(i), 32, 64)
             for i in range(3)]
        ),
        learning_rate_init=trial.suggest_loguniform("lr_init", 1e-5, 1e-1),
    )

    for step in range(100):
        clf.partial_fit(x_train, y_train, classes=classes)
        accuracy = clf.score(x_valid, y_valid)
        error = 1 - accuracy

        # 中間評価値を報告.
        trial.report(error, step)

        # 早期停止すべき状態なら TrialPruned 例外を送出.
        if trial.should_prune():
            raise minituna.TrialPruned()
    return error

枝刈りを利用する際に必要なAPIは2つだけです。

1. trial.report(value, step) で中間評価値をStorageに保存
2. trial.should_prune() が True なら、TrialPrued() 例外を送出し、学習を停止

Optuna Trial APIの設計 (Optunaの論文 [1] Figure.6 より参照)

このAPIからわかるように現状Optunaがサポートしている枝刈りアルゴリズムは、すべて中間評価値の情報をもとに枝刈りの判断を行います。MinitunaはMedian stopping rule [2] というアルゴリズムを実装していますが、幸いにもOptunaがサポートしている枝刈りのアルゴリズムは、どれもそこまで複雑なものではありません。

これらのアルゴリズムはすべて、中間評価値が低ければ、最終評価値もそんなによくないだろうという経験的な知見をもとに途中で学習を停止させます。このあと紹介するMedian stopping ruleの他に、OptunaではSuccessive Halving [3, 4] やHyperband[5]といったアルゴリズムもサポートしていますが、基本的な考え方は大きくは変わりません。弊社AI Lab所属の研究員野村のPyCon JPでの発表 「機械学習におけるハイパーパラメータ最適化の理論と実践」で枝刈りのアルゴリズムも解説されているので気になる方はチェックしてみてください。

AI Lab 研究員 野村のPyCon JP 2019発表資料より抜粋
https://youtu.be/F1GGPQlra-E?t=1617

Median stopping ruleでは基本的に各stepで下位半分(過去のtrialの中間評価値のMedian値より悪いもの)を枝刈りします。Median stopping ruleの挙動を図示すると次のようになります。コードもシンプルなので次の図と合わせて読んでみるとMedian stopping ruleアルゴリズムの詳細がよくわかるでしょう。

Median stopping ruleの挙動

OptunaのSuccessiveHalvingやHyperbandの実装を読んでみたい方には、1つ注意点があります。Optunaは特定のTrialの学習を中断/再開できるような設計にはなっていないため、アルゴリズムに少し改変を加えているため、論文で解説されているアルゴリズムとは異なった挙動をします。Optunaが実装しているSuccessive HalvingについてはOptunaの論文 [1] のAlgorithm 1、Hyperbandについては実装者のcrcrparさんによる解説ブログ「How We Implement Hyperband in Optuna」を参照してください。

Prunerの設計に関する課題感についても触れておきます。minitunaのコードからわかるようにOptunaではPrunerとSamplerのinterfaceは明確に分離されています。そのおかげでOptunaのコードは見通しがよく、SamplerとPrunerを自由に切り替えて組み合わせることができます。これは大きなメリットです。

その一方でPrunerとSamplerが密に連携する必要があるアルゴリズムもあり、それらは今の設計では実装できません。Hyperbandの実装においても実は裏側で少しトリッキーな実装を加えることでPrunerとSamplerを連携させて動かしています。PrunerやSamplerのインターフェイスには未だ議論の余地があります。もしこのあたりのコードを読んだうえで、何かアイデアがある方は提案をいただけると嬉しいです。

Define-by-RunインターフェイスにおけるRelative Samplingの仕組み

Minitunaでは実装ボリュームの都合上、解説することができなかったRelative Samplingについても触れておきます。このあたりはDefine-by-Runインターフェイスを採用したOptuna特有の仕組みです。SkoptSampler (GP-BO)やCmaEsSamplerのようにパラメーター間の相互作用を考慮できる最適化アルゴリズムを実装する場合には、Relative Samplingの理解が必要になります。ベイズ最適化や進化戦略アルゴリズムの研究をされている方は、Relative Samplingまで理解すれば、独自のSamplerをOptuna上で実装できるでしょう。

MinitunaのSamplerインターフェイスの設計は、Optuna v0.12.0時代のインターフェイスとほぼ同じものです。しかしOptuna v0.13.0以降のインターフェイスは、これらと異なります。次の2つのSamplerインターフェイスを見比べるとその違いがわかります。

• v0.12.0: https://github.com/optuna/optuna/blob/v0.12.0/optuna/samplers/base.py
• v2.0.0: https://github.com/optuna/optuna/blob/v2.0.0/optuna/samplers/_base.py

Relative Samplingを理解するために、いくつかの目的関数を例にそれぞれの探索空間がどのようになっているか考えてみましょう。

def objective(trial):
    x = trial.suggest_uniform("x", -100, 100)
    y = trial.suggest_categorical("y", [-1, 0, 1])
    return x ** 2 + y

この目的関数の探索空間は常に↓のようになります。

{
    'x': UniformDistribution(low=-100, high=100),
    'y': CategoricalDistribution(choices=[-1, 0, 1])
}

では次の目的関数の探索空間はどうでしょうか？

def objective(trial):
    x, y = ...

    classifier_name = trial.suggest_categorical("classifier", ["SVC", "RandomForest"])
    if classifier_name == "SVC":
        svc_c = trial.suggest_float("svc_c", 1e-10, 1e10, log=True)
        classifier_obj = sklearn.svm.SVC(...")
    else:
        rf_max_depth = trial.suggest_int("rf_max_depth", 2, 32, log=True)
        classifier_obj = sklearn.ensemble.RandomForestClassifier(...)

    accuracy = ...
    return accuracy

Define-by-Runインターフェイスでは、実行時に探索空間が決定します。この例では、探索空間にif文による条件分岐が存在するため、探索空間は1つではありません。次の2つの探索空間がコンテキストに応じて切り替わります。

# SVC
{
  'classifier': CategoricalDistribution(choices=["SVC", "RandomForest"),
  'svc_c': LogUniformDistribution(low=1e-10, high=1e10),
}

# RandomForest
{
  'classifier': CategoricalDistribution(choices=["SVC", "RandomForest"),
  'rf_max_depth': IntUniformDistribution(low=2, high=32),
}

そこでCmaEsSamplerやSkoptSamplerはすべての探索空間で共通して登場しているものを、 Sampler.infer_relative_search_space(study, trial) メソッドで取り出し、 Sampler.sample_relative(study, trial, search_space) の第3引数に渡します。つまり上の例では classifier パラメーターのみがRelative Search Spaceとして扱われます。このRelative Search SpaceからのSampleにのみGP-BOやCMA-ESを使っていて、これをRelative Samplingと呼んでいます。一連の流れを図にしたものがこちらです。

Samplerの呼び出しフロー (公式ドキュメンテーションより参照)

svc_c や rf_max_depth のようにRelative Search Spaceに含まれなかったパラメーターは、RandomSamplerやTPESamplerのように変数間の依存関係を考慮しない手法を実装したSamplerにfallbackします。

まとめ

本記事ではMinitunaの紹介とMinituna/Optunaのコードを読みすすめる上での補足解説を行いました。Minituna v2まで写経し、全体の流れがつかめた方であればOptunaのコードは読めると思います。興味のあるコンポーネントから読み進めてみてください。Optuna開発チームのみなさんは非常にsupportiveで、PRも丁寧にレビューしてくれます。この記事をきっかけにぜひOptunaの開発に参加してみてください。

参考文献

1. T. Akiba, S. Sano, T. Yanase, T. Ohta, and M. Koyama:  Optuna: A Next-generation Hyperparameter Optimization Framework. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 2623-2631, 2019.
2. D. Golovin, B. Sonik, S. Moitra, G. Kochanski, J. Karro, and D.Sculley: Google Vizier: A Service for Black-Box Optimization. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 1487-1495, 2017.
3. K. Jamieson and A.Talwalkar: Non-stochastic Best Arm Identification and Hyperparameter Optimization. In Proceedings of the 19th International Conference on Artificial Intelligence and Statistics (AISTATS),  PMLR 51:240-248, 2016.
4. L. Li, K. Jamieson, A. Rostamizadeh, E. Gonina, M. Hardt, B. Recht, and A. Talwalkar: A System for Massively Parallel Hyperparameter Tuning. arXiv preprint arXiv:1810.05934, 2018.
5. L. Li, K. Jamieson, G. DeSalvo, A. Rostamizadeh and A. Talwalkar: Hyperband: A Novel Bandit-Based Approach to Hyperparameter Optimization. Journal of Machine Learning Research, 18, pp. 1-52, 2017.

謝辞

本記事はOptunaの開発チームのみなさんやスキルアップAI株式会社の斉藤さんにレビューしていただきました。この場でお礼をさせていただきます。ありがとうございました。

CMA-ESをPythonで実装しGitHubで公開したのでその使い方や性能について紹介します。Optunaからも利用できるようになっているのでぜひハイパーパラメータの最適化に使ってみてください。

URL: cmaes: Lightweight Covariance Matrix Adaptation Evolution Strategy (CMA-ES) implementation for Python.

CMA-ES

CMA-ESは多変量正規分布から解を生成し、その解の評価値を利用して、より良い解を生成するような分布に更新を行う手法です。変数間の依存関係のある問題や変数間のスケールが異なるような問題に対して有力であることが報告されています [3]。また、ある程度の評価回数を許容できる場合に、様々なブラックボックス最適化問題に対する実験において、100以上の最適化手法の中で最も良い性能を示したことが Loshchilovらによって報告されています[4]。

PythonライブラリとしてはCMA-ESの提案者Nikolaus Hansenが公開しているpycmaが最も有名です。Optunaにもpycmaとのintegration samplerが実装されています。しかしpycmaの設計は弊社の野村がNeurIPS 2019のMeta Learning workshopで発表した手法を実装する際にインターフェース上の問題があり、コード自体も複雑になってきてしまっている面があるため、今後のメンテナンス性を踏まえフロムスクラッチで実装しました。

CMA-ESのアルゴリズムについては、The CMA Evolution Strategy: A Tutorial  [5]を参照してください。ユニットテスト等による実装の確認が難しいことから、matplotlib.animationベースの可視化ツールを実装しこちらで挙動を確認しながら実装を進めていきました。多変量正規分布の共分散行列が最適解付近で小さくなっていく様子が確認できます。

次のように変数間のスケールが異なる問題でも等高線に沿うような形で共分散行列が更新されていることが確認できます。

cmaesの使い方

それではライブラリの使い方を簡単に紹介します。PyPIにも上げているので、pipからインストールできます。

$ pip install cmaes

このライブラリは2種類のインターフェースを提供しています。OptunaのSamplerとしてつかえるインターフェース (cmaes.sampler.CMASampler)と、pycmaのように ask() と tell() メソッドを提供する低レベルなインターフェース(cmaes.CMA) です。基本的には cmaes.sampler.CMASampler を用いてOptunaから利用することがおすすめです。Optuna samplerを使う場合は、次のようになります。

import optuna
from cmaes.sampler import CMASampler


def objective(trial: optuna.Trial):
    x1 = trial.suggest_uniform("x1", -4, 4)
    x2 = trial.suggest_uniform("x2", -4, 4)
    return (x1 - 3) ** 2 + (10 * (x2 + 2)) ** 2


def main():
    sampler = CMASampler()
    study = optuna.create_study(sampler=sampler)
    study.optimize(objective, n_trials=250)


if __name__ == "__main__":
    main()

注意点として、Categorical distributionは現状サポートしていません。Optunaが提供するpycmaベースのsamplerは、質的変数をサポートしていますが、label encodingにより連続空間に落としているため、カーディナリティの高い質的パラメーターが含まれる場合にはうまく最適化できないケースがあるかと思います。

そのためCMASamplerは、 trial.suggest_categorical を使っておらず、n_trialsを十分大きくできるときに利用してみてくださいい。trial.suggest_categoricalを使っている場合は基本的にデフォルトのTPE samplerを利用するのがおすすめです。

kurobakoとGitHub Actionsを使った継続的ベンチマーク

CMA-ESの実装の妥当性をユニットテスト等で検証するのが難しいことから、アルゴリズムの性能を確認できるベンチマークがあると非常に有用です。ベンチマークには  ブラックボックス最適化ベンチマークツール「kurobako」の紹介 で解説されているkurobakoを利用しました。

ベンチマーク関数はRosenbrock, Six-Hemp Camel, Himmelblauの3種類を用意しました。今回公開したCMA-ESの実装はアルゴリズム自体のコード行数が200行程度と比較的シンプルにおさまりましたが、その性能はpycmaと同程度であることが確認できています。

GitHub Actionsでの実行にはgithub-actioins-kurobakoというツールを作成・使用しました。Pull Requestを作成するとGitHub Actionsがkurobakoを実行し、その結果をPRにコメントすることで、ベンチマーク結果を継続的に追えるようにしています。諸事情でCyberAgent organizationはGitHub Actionsが使えないため自分の個人アカウントにforkしてこちらのPRのように動かして確認しています。

高速なOptuna sampler実装

CMA-ESの魅力的な点の1つは、次に探索すべきパラメーターの生成に必要な計算量が過去の試行回数に依存しない点です。GPやTPEでは代理モデルの構築時に過去の試行結果をすべて取り出してくる必要がありますが、CMA-ESで必要になるのは基本的に多変量正規分布の平均ベクトルと共分散行列の値だけです。

その一方でこちらのOptunaのIssueにあるようにOptunaが想定しているユースケースである分散最適化等を考慮すると、実装上悩ましい点がいくつかあります。現状、Optunaが提供しているpycmaのintegrationは試行回数が積み上がるにつれ必要な計算量が増加していく課題があります。

今回公開している実装では、多変量正規分布の平均ベクトルや共分散行列等をstorage上でキャッシュし、分布の世代番号(更新されるたびにインクリメントされる値)をチェックしながら最新分のtrialを追加適用しています。またOptunaの内部の関数も一部再実装しました。高速化についてはまたOptuna側にもPRを送っていけたらと思います。

おわりに

本記事ではOptunaから利用可能なCMA-ES実装の解説を行いました。シンプルな実装ながらpycmaと同程度の性能であることを確認し、Optunaから利用する際には比較的高速に動作するように実装しています。みなさんもぜひ使ってみてください。

参考文献

[1] Friedrichs, F. and Igel, C. Evolutionary Tuning of Multiple SVM Parameters.Neurocomputing, 64:107–117, 2005.
[2] Watanabe, S. and Le Roux, J. Black box optimization for au-tomatic speech recognition. In 2014 IEEE InternationalConference on Acoustics, Speech and Signal Processing (ICASSP), pp. 3256–3260, 2014.
[3] Rios, L. M. and Sahinidis, N. V. Derivative-free optimization: a review of algorithms and comparison of software implementations. Journal of Global Optimization, 56(3):1247–1293, 2013.
[4] Loshchilov, I., Schoenauer, M., and Sebag, M. Bi-population CMA-ES Algorithms with Surrogate Models and Line Searches. In Proceedings of the 15th annual conference companion on Genetic and evolutionary computation, pp. 1177–1184, 2013.
[5] N. Hansen. The CMA Evolution Strategy: A Tutorial. arXiv:1604.00772, 2016.
[6] Loshchilov, I. and Hutter, F. CMA-ES for Hyperparameter Optimization of Deep Neural Networks. In ICLR Workshop, 2016.