🛠️

MLflowで画像認識案件の実験管理を効率化する

2022/10/21に公開

こんにちは。SREホールディングス株式会社の西野です。
主に画像認識案件のPLを担当しています。
本日は画像認識チームにおいて実験管理目的で採用したMLflowについてご紹介いたします。

背景

弊社では不動産に限らず、様々な業界・分野におけるAIの社会実装を推進しており、最近では画像認識案件にも注力するようになっています。
そのような中、私はPLという立場で複数のアルゴリズムエンジニアと連携しながら複数のプロジェクトを推進していくわけですが、その中で課題感としてあったのが実験管理の効率化です。

  • クライアントに報告したこの精度結果って、どの学習済みモデルでしたっけ?
    • そのときのハイパーパラメータは?
    • どのデータセットで?
  • この実験ってどういう意図で実施したんでしたっけ?

Excelなどを使って管理しようとはしても、煩雑でミスも起こりやすいため、
(今後のMLOps的な文脈の期待も込めて)MLflowを導入して解決を図りました。

対象読者

  • MLflowに興味がある方

概要

本記事ではMLflowの実験管理機能について、

  • MLflowでできること
  • 弊社における環境構築

などについて簡単にご紹介いたします。
(詳しい環境構築方法等については記載していません)

MLflow でできること

まずは「MLflow の実験管理(MLflow Tracking)でできること」について、
実際の画面ベースで簡単に紹介します。

webUIで実験結果の一元管理

こちらが実際の実験管理の画面になります。Webブラウザでアクセスする形になります。
Webアプリなので、閲覧するだけであれば特に環境構築の必要が無く、簡単に情報にアクセスできるところが、(特にPLの立場の私にとっては)魅力的です。

概念としては、実験、RUN、各種ログというシンプルな階層になっています。

各RUNに記録できるログ

こちらがひとつのRUNを選択して表示した画面です。以降、上から順に説明します。

自動保存

以下の情報はMLflowが気を利かせて自動的に記録してくれます(MLflow system tags)。
git commitは有難いですね。

  • RUNの実行日時
  • ソースファイル名
  • git commitハッシュ値(git管理している場合)
  • RUNを実施したユーザー
  • RUNの実行時間

など

Description

任意の文字列が記録できます。"概要"、"説明"ですので、このRUNを実行した背景や狙いなどを記録するのが一般的と思われます。
WebUI上で後から編集可能です。

Parameters

キーバリュー形式で任意の文字列を記録できます。
ハイパーパラメータなどを記録します。

Metrics

同じくキーバリュー形式ですが、値が数値であるメトリクス(評価指標)を記録します。
モデルの損失関数などを記録するのに利用します。
以下のように自動的に可視化してくれます。

Artifacts

任意のフォーマットのファイルを記録できます。
学習済みモデルやデータセットなどを保存できます。
保存したファイルは、WebUI上からダウンロード可能です。(←これすごく便利です)

ロギング方法

Tracking Serverが構築できていれば、クライアント側(推論/学習コード側)は、コードに少し記述を加えるだけで簡単にログを記録ができます。

Tracking Serverの指定

まず初めに、トラッキングサーバーを指定し、実験IDを指定します。

import mlflow

TRACKING_URI = 'http://(トラッキングサーバーのIPアドレス):5000'
mlflow.set_tracking_uri(TRACKING_URI)
EXPERIMENT_NAME = '(任意の実験ID)'
experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)
if experiment is None:  
    experiment_id = mlflow.create_experiment(
                            name=EXPERIMENT_NAME)
else: 
    experiment_id = experiment.experiment_id

各種ログの保存

あとはmlflow.start_run()でRUNを開始し、
mlflow.log_parammlflow.log_metricmlflow.log_artifactで適宜ログを保存するだけです。

with mlflow.start_run():
    mlflow.log_param("x", 1)

    for epoch in range(0, 3):
        mlflow.log_metric(key="quality", value=2*epoch, step=epoch)

    mlflow.log_artifact("features.txt")

※Artifactの保存先については、(今回弊社が作成した構成では)クライアント側は特に意識する必要はありません。
弊社ではArtifactの保存はFTPサーバを採用しましたが、環境構築時点で、Tracking サーバに対して、利用するArtifactのFTPサーバを指定していて、FTPのやりとりはTrackingサーバが勝手にやってくれるので、クライアント側は意識する必要はありません。

弊社における環境構築

ここからは主に環境構築の話に移りますが、詳しい環境構築方法等については、下記の記事等に良くまとまっていますので、本記事では簡単な紹介程度にとどめます。

MLflowの環境構築方法まとめ【Docker, S3】

要件整理

環境構築についてはMLflowの公式ドキュメントにおいても、6つのシナリオが紹介されている通り、構成についてはいくつか選択肢があり、要件にあわせて選択できます。
(結果的に弊社に構築した環境はシナリオ5の構成に近いです)
例えば、Artifactの保存(ストレージの選択)については以下のような選択肢があります。

  • ローカルストレージ / FTPサーバ / S3

弊社では色々な諸条件を加味したうえで、画像認識系の学習環境はクラウドではなくオンプレに構築しています。
社内のサーバールームにGPUマシンが数台設置されていて、そこに各エンジニアがリモートでアクセスして学習や推論などの検証を実施する形になります。
したがって、(環境構築観点では)以下のような要件があります。

  • オンプレに環境を構築する(画像データはオンプレに保存されている)
  • 複数のエンジニア(GPUマシン)による実験を一元管理できる

構築した環境

これを実現するにはMLflowのTrackingサーバーを活用します。
Trackingサーバーを用いることで、クライアントはRESTのAPI経由でリモートからもMLflowを利用できます。
(といってもクライアント側でRESTAPIを意識する必要はありません)

弊社では以下の図のような構成で環境を構築しています。

まとめ

本記事ではMLOpsのOSSツールであるMLflowを簡単にご紹介させて頂きました。
個人的には(PLの立場として)以下の面で業務の効率化が図れたと感じています。

  • アルゴリズムエンジニアとのコミュニケーション
  • 再現性の確保
  • クライアントへのリリース管理

今後も様々な便利ツールを駆使しながら、画像認識の社会実装を実現していきたいと思います!

参考

SRE Holdings 株式会社

Discussion