Optunaの最適化クラスを作った話

株式会社ブレインパッドでデータサイエンティストをしている23新卒のまーしです。

この記事はBrainPad Advent Calender 2023 17日目の記事です。

本記事の目的

Optunaをシステムに組み込むために作った最適化クラスのコード共有です。

はじめに

みなさんは、Optunaを使ったことはあるでしょうか？
Optunaとは、日本のPrefferdNetworks社が開発したハイパーパラメータ（ハイパラ）の自動最適化フレームワークです。より広義にはハイパラ最適化を含むブラックボックス最適化全般に使用可能なフレームワークです。

業務でハイパラ最適化でないブラックボックス最適化でOptunaを用いる必要がありました。その際に工夫した点が次の2つです。

1. 制約遵守解のみを出力
2. 最適化クラスを設計

工夫点

1. 制約遵守解のみを出力

Optunaには制約付きブラックボックス最適化ソルバーがいくつか実装されています。ただし、どのソルバーも制約条件をソフト制約として扱います。つまり、探索で得られる解には制約違反解と制約遵守解両方が混在します。業務上、制約遵守解のみ出力したかったため、以下のような制約遵守解だけ取り出す関数¹を作成しました。

def get_feasible_trials(self):
    # 全ての試行のうち、制約条件を全て満たすものだけを抽出し、目的関数の値でソート(昇順)
    feasible_trials = sorted(
        [t for t in self._study.trials if all(np.array(t.user_attrs["constraints"]) <= 0)],
        key=lambda t: t.value,
    )

    return feasible_trials

2. 最適化クラスを設計

機械学習モデルをクラスとして定義するように、Optunaもクラスとして定義することで、よりシステムに組み込みやすくしました。目的関数や制約条件とクラスの実装は以下の通りです（まだまだ改良点があるのでコメント頂けたら幸いです）

import optuna
import numpy as np

#目的関数1
def f0(X):
    return 4 * X[0] ** 2 + 4 * X[1] ** 2
#目的関数2
def f1(X):
    return (X[0] - 5) ** 2 + (X[1] - 5) ** 2

# 制約条件1
def g0(X):
    return (X[0] - 5) ** 2 + X[1] ** 2 - 25

# 制約条件2
def g1(X):
    return -((X[0] - 8) ** 2) - (X[1] + 3) ** 2 + 7.7

class OptunaModel:
  def __init__(self):
    super(OptunaModel, self).__init__()
    self._sampler = optuna.samplers.TPESampler(
        constraints_func=lambda x: x.user_attrs["constraints"],
    )

    self._study = optuna.create_study(
        directions=["minimize"],
        sampler=self._sampler,
    )

  def _add_variables(self,trial):
    self._X = [trial.suggest_int(f"x{i}", -15.0, 30.0) for i in range(2)]


  def _add_constraints(self, X, trial):
    C = []
    for const_func in [g0, g1]:
        C.append(const_func(X))
    trial.set_user_attr("constraints", C)

  # 目的関数の重み付け和をとり、単目的最適化にする
  def _get_objective(self, X, trial):
    total_obj = 0
    for i, obj_func in enumerate([f0, f1]):
        total_obj += 10**(-i) * obj_func(X)
    return total_obj

  def _objective(self, trial):
      # 決定変数の作成
      self._add_variables(trial)
      # 制約条件の作成
      self._add_constraints(self._X, trial)
      # 目的関数の作成
      obj = self._get_objective(self._X, trial)
      return obj

  def optimize(self):
    self._study.optimize(self._objective,timeout=2000)

上記のクラスは次のようにして使用できます。

model = OptunaModel()
model.optimize()
feasible_trials = model.get_feasible_trials()
# 制約条件を満たす試行のうち最良の目的関数値と解を表示
print(feasible_trials[0].value)
print(feasible_trials[0].params)

まとめ

Optunaをハイパラ最適化以外のブラックボックス最適化に適用する場合、「制約遵守解のみ欲しい」や「システムに組み込める実装にしてほしい」の要望が出るかもしれません。その際、上記のコードがわずかながらでも参考になれば幸いです。

1. 単目的の最小化問題を前提としたコーディングになっています。 ↩