Kubeflow/KatibがGoptunaを使った最適化に対応しました。
KubeflowのKatibというハイパーパラメーター最適化等を担当するコンポーネントにGoptunaを使ったSuggestion serviceを実装しました。
Goptunaはハイパーパラメーター最適化ライブラリとして、機能面でも実装品質の面でもPythonで人気のライブラリに劣らないものになってきたかなと思う一方で、Goの機械学習ユーザーはまだまだ少なく、その中でも学習までGoでこなしている人はさらに絞られるので、せめてこういった方面で使われていくと嬉しいなと思います。ブラックボックス最適化フレームワークとしていろんな用途に利用できるソフトウェアなのでみなさんもぜひ触ってみてください。
Katibの基本的な使い方
Katibは基本的に特定の言語やフレームワークに依存しないように設計されています。パラメーターや評価値の受け渡しをどうやっているかというと、コマンドライン引数などからパラメーターを受け取り、目的関数の評価値を予め決めておいたフォーマットで標準出力に出したり、ログファイルとして保存したりします。例として f(x1, x2) = (x1-5)^2 + (x2+5)^2
を目的関数として用意しました (GitHub: https://github.com/c-bata/katib-goptuna-example)。
import argparse import logging logging.basicConfig(filename='/var/log/katib.log', level=logging.DEBUG) def main(): parser = argparse.ArgumentParser() parser.add_argument("--x1", dest='x1', type=float) parser.add_argument("--x2", dest='x2', type=float) args = parser.parse_args() evaluation = (args.x1 - 5) ** 2 + (args.x2 + 5) ** 2 logging.info(f'{{metricName: evaluation, metricValue: {evaluation:.4f}}};') if __name__ == '__main__': main()
FROM python:3-alpine ADD main.py /usr/src/main.py WORKDIR /usr/src
この関数を最適化する際にExperiment CRを作成します。OptunaでいうところのStudyに対応するものです (昔はStudyと呼ばれていたのか、コード読んでるとStudyって書かれてるところもありますが、詳細はまだよく分かっていません)。
Experiment CRが作成されると、パラメーターを変えながらKubernetes Jobを何個も作成し先程のPythonファイルを実行します。各パラメータによる試行はTrialと呼ばれます。Experiment CRで指定するのはざっくり次のような項目です。
- 目的関数の評価回数や並列数など
- 目的関数に与えるパラメーターの探索範囲
- 探索アルゴリズム (文字列で
tpe
やcmaes
を指定)- 各アルゴリズムがどのSuggestion serviceに対応するかは、
katib-config.yaml
で指定します。 - デフォルトは https://github.com/kubeflow/katib/blob/master/manifests/v1alpha3/katib-controller/katib-controller.yaml です。
- 質的変数を探索空間に含む場合には
tpe
を利用する方が多いかと思いますが、デフォルトではhyperopt
ベースのsuggestion serviceが利用されます。Goptunaにもtpe
が実装されているので、対応するイメージをGoptuna suggestion serviceのイメージに書き換えればCMA-ESだけでなくTPEでもGoptunaを利用します。
- 各アルゴリズムがどのSuggestion serviceに対応するかは、
- metricscollectorがどのファイルからどういう形式で評価値を取り出すか
- Kubernetes Jobのmanifest template。どういうふうにコマンドライン引数からパラメーターを与えるかはここで指定できます。
具体例は次のとおりです。
apiVersion: "kubeflow.org/v1alpha3" kind: Experiment metadata: namespace: kubeflow labels: controller-tools.k8s.io: "1.0" name: example spec: objective: type: minimize goal: 0.001 objectiveMetricName: evaluation algorithm: algorithmName: cmaes metricsCollectorSpec: source: filter: metricsFormat: - "{metricName: ([\\w|-]+), metricValue: ((-?\\d+)(\\.\\d+)?)}" fileSystemPath: path: "/var/log/katib.log" kind: File collector: kind: File parallelTrialCount: 2 maxTrialCount: 250 maxFailedTrialCount: 3 parameters: - name: x1 parameterType: double feasibleSpace: min: "-10" max: "10" - name: x2 parameterType: double feasibleSpace: min: "-10" max: "10" trialTemplate: goTemplate: rawTemplate: |- apiVersion: batch/v1 kind: Job metadata: name: {{.Trial}} namespace: {{.NameSpace}} spec: template: spec: containers: - name: {{.Trial}} image: docker.io/cbata/hello-katib-quadratic-function command: - "python3" - "main.py" {{- with .HyperParameters}} {{- range .}} - "--{{.Name}} {{.Value}}" {{- end}} {{- end}} restartPolicy: Never
準備に必要なファイルはこれだけです。あとは Katib リポジトリにある、 scripts/v1alpha3/deploy.sh
を実行してKatib controller等を立ち上げ、上記のExperiment CRDを適用すれば最適化が実行されます。
試行結果は kubectl port-forward svc/katib-ui -n kubeflow 8080:80
を実行してKatib UI https://localhost:8080/katib をブラウザで開いて確認できます。
Suggestion serviceの実装
Suggestion serviceは、過去のtrialsを受け取り次に探索すべきパラメーターを返すgRPCのサーバーです。 基本的にKatibを利用する上では知る必要はそんなにないので、適宜読み飛ばしてください。 Katibの動作の流れは次の画像にまとまっています。
katib/suggestion.md at master · kubeflow/katib · GitHub
Experimentの作成後、Experimentで指定されたalgorithmNameをもとに katib-config
から対応するSuggestion serviceのImageを特定し、Katib controllerがSuggestion CRをapply、Suggestion CRがSuggestion serviceのdeploymentをapplyしPodが生成されます。
Suggestion serviceがreadyになったあと、Katib controllerは過去の試行結果を入力にして次に探索するパラメーターを取得する GetSuggestions()
gRPC エンドポイントを何度も叩いてきます。Suggestion controllerはsuggestion serviceを1台しか立ち上げないため、各workerが勝手に解をサンプルしていくOptunaやGoptunaと比べアルゴリズムは実装しやすいなと思いました。Suggestion serviceを実装する際に問題になりやすいのは、Katib controllerがExperimentに紐づけて保持する過去の試行結果と、Goptunaの内部状態をどうやって同期するかです。
Katibの設計ではSuggestion serviceはただパラメーターをサンプルし、Katib controllerに伝えます。その後Katib controllerがtrialを生成し、一意なIDであるtrial nameを生成します。Suggestion serviceからみたときに、自分が生成したパラメーターにどのtrial nameが紐づくかを特定する方法は基本的にありません。TPEやRandom search、Gaussian Processなどこれまでサポートされてきた多くの最適化手法ではこれは問題になりませんが、CMA-ESは別です。
CMA-ESは、多変量正規分布を用意しその分布から解を生成、その評価値をもとに多変量正規分布のパラメーターをよりよい解をサンプルする分布へと更新します。分布が更新されるごとに世代番号がインクリメントされ、分布の更新に利用する解はかならず同じ世代番号の分布からサンプルされている必要があります。分布の世代番号はGoptunaのTrialのメタデータとして紐付けているため、Goptuna TrialとKatib Trialを適切に紐付ける必要があります。
現状ではGoptunaのStudyとは別に trialMapping map[str]int
という変数を用意して、Katibのtrial name (str)からGoptunaのtrial ID (int)へのmappingを用意し、 trialMapping
にまだ紐付けされていないtrialの中からパラメーターが完全に一致したものをtrialMappingへ追加しています。初期の実装ではKatib-controllerとのやりとりの際に生じる桁落ちを懸念して、パラメーターどうしのマンハッタン距離から最も似ているtrialを紐付けていましたが、Katibはprotobufの定義も(おそらくetcdかなにかに保存するときも)内部表現は全て文字列です。そのため桁落ちなどの心配がなく完全一致でチェックすることにしました。あらためて考えるとTrialのパラメーター表現を文字列で統一するのはかなりわかりやすくリーズナブルだなと思います。
あと実装してから気づいたのですが、よく見たらSuggestion serviceを実装するためのドキュメントがありました。普段あまりドキュメントに期待していないのですが、ParameterAssignmentsが requestNumber
で指定された個数だけパラメーターを返さないといけないこととか気づくのに1時間ぐらい時間溶かしたりしてたので、最初に読んでおくことがおすすめです。
katib/new-algorithm-service.md at master · kubeflow/katib · GitHub
所感と今後の展望
設計的にも面白いソフトウェアだなと思いました。また今はOptunaのcommitterもやっていますが、前の部署にいたときにGoでgRPCのサーバー書いたり、OSSでkube-promptを公開していたりしていたので、スキルセット的にも結構マッチしている気がして開発が楽しかったです。
一方で、運用できるチームはかなり限られるかもしれないなというのが正直な感想です。少し動かしただけでもhyperoptのsuggestion serviceが完全に動かなくなってしまっているデグレがあって、原因探って修正するまでに自分も1晩中頑張ったりして苦労しました。Suggestion serviceがおかしいのか、metricscontrollerがおかしいのか、Katib controllerがおかしいのかを正しく状況に応じて切り分ける必要があり、Katibのコード読んでる人でもない限りなかなか難しいかもしれません。
個人的にはGoptunaをベースにすればKatibとの互換性を保ったままシンプルで運用のしやすいツールができる気もしているので、kubebuilder触りつつまた気が向いたときにでもやってみたいなと思います (NASの対応は、言語やフレームワークに依存せず汎用的に使える便利なツールを設計するのが現状では難しい気がしてるので諦める予定です)。