MLFlowと他ツールの組み合わせ

2023年10月11日 12:41

※こちらの記事は、2020年7月28日にRetrieva TECH BLOGにて掲載された記事を再掲載したものとなります。

こんにちは。カスタマーサクセス部リサーチャーの坂田です。
レトリバでは、固有表現抽出、分類、PoC用ツール作成に取り組んでいます。
PoC用ツール作成は、研究成果をより迅速にPoCで試せることを狙いとしています。実験結果の可視化UIが充実しているMLFlow を中心に、足りないところを補うため、その他のツールとの組み合わせについて考えていきます。

MLFlow

MLFlow は、実験管理からデプロイまでカバーしたツールです。特定のツールに依存しないということに重きを置いています。
4つのコンポーネントに分かれており、必要な機能のみを使えるようになっています。

MLflow Tracking : パラメータ、コードのバージョン管理、生成物の捕捉などを行う機能など。
MLflow Projects : 再現性を担保するための機能など。
MLflow Models : デプロイの支援機能など。
MLflow Registry : モデルのバージョン、ステータス(例：Staging、Production、Archived)を管理など。

MLFlowでは、実験結果が以下の2つの単位で管理されます。

Experiment : 比較を行いたいRunをまとめた概念。
Run : パラメータとそのパラメータに紐づいたスコアをまとめた概念。各Runは、1つのExperimentに属する。

UIの左側の一覧からExperiment を選択できます。Experiment名での検索も可能です。選択したExperimentに属するRunがUIの右側に表示されます。こちらは、パラメータがある特定の値のとき、スコアがある値以上などの条件を指定した検索が可能です。

Runを選択し、Compareボタンを押すと該当Runを比較するUIが表示されます。

ExperimentごとにMarkdown形式で説明を記述できるのも便利です。

可視化UIを使いたいだけならば、MLFlow Trackingのみを使って以下のように簡単に使えます。

mlflow.set_experiment(‘Test’)
with mlflow.start_run(run_id=5):
    mlflow.log_param('param1', 0.5)
    mlflow.log_param('param2', 0.3)
    # 何かしらの機械学習モデルの学習をする。
    mlflow.log_metric('score', 0.88)

上記の例では明示的にRun IDを5と指定しています。動かすPythonスクリプトが1つだけならば特に指定する必要はありませんが、前処理、訓練、評価など複数スクリプトを分けた場合に、IDが一致しないと同一Runとして扱われません。UI上でもバラバラに表示されてしまいます。他ツールとの組みあわせの際にもこのことが障害になる場合があります。

その他のツール

以下では、MLFlow と組み合わせることを前提に、2つのツールをご紹介します。 Kedroはパイプラインの可視化、Hydraは複数パラメータでの実行と階層性を持ったパラメータ指定が出来るメリットがあります。

Kedro

Kedroは、下図のように機械学習パイプラインの可視化UIが充実しています。

Kedroは、以下の2つを重視しています。PR動画でもこれらがアピールされています。

モジュール性の担保する。
テンプレートを埋めることで製品レベルのコードが書けるようにする。

各機能をモジュール化することによって、運用環境で処理が止まった時にどの部分に原因があったのか明確になり、一部機能の切り替えが容易になります。また、テンプレートによって、各人のコードの書き方の差異が強制的に減少するため、お互いにコードを理解しやすくなります。

KedroにはData Catalogという機能があります。訓練データ、テストデータに加えて、各ノードの中間生成物も以下のようにYAML形式で記述しておきます。

そして、前処理、訓練、評価といった機械学習パイプラインの各ノード間のやりとりは、Data Catalogに定義された名前を渡すことで行います。これによって、様々な形式のデータの読み込みの実装が分離され、コードが読みやすくなります。また、全ての入出力データが一覧出来るメリットもあります。

一方で、他の様々なツールと組み合わせるという点では、以下のように、多少煩雑となる場合があります。

ノード間で単純な数値・文字列を中間生成物渡す際に、中間生成物を介する必要がある。(MLFlow の Run IDも含まれる。）
パラメータをYAML形式で指定するため、次に紹介するHydraとの併用にはひと工夫が必要。

Hydra

複数のパラメータで実行することを支援するツールです。以下のようにパラメータそれぞれを複数指定すれば、グリッドサーチを行ってくれます。この例では、(mecab-dic 3通り) × (embedding 2通り) = 6通りのパラメータで実行されます。

python main.py mecab-dic=ipadic,neologd,unidic embedding=word2vec,fasttext -m

パラメータの初期値をYAMLに記述しておき、変更したいパラメータのみをコマンドラインから指定する仕様になっています。 argparseやclickのようなコマンドラインツールを兼ねており、デコレータによって実装します。 MLFlow と組み合わせる際の実装も、パラメータを渡すところのみを変更すればいいため比較的容易です。以下に、MLFlow 単体にclickという別のコマンドラインツールでパラメータを渡した例と、hydraを使った例を示します。書き換える量はそれほど多くありません。

~略~
@click.command()
@click.option('--mecab-dic', default='ipadic', type=str)
@click.option('--embedding', default=None, type=str)
def workflow(data, mecab_dic, post_process):
    client = mlflow.tracking.MlflowClient()
    experiment_list = client.list_experiments()

    # データごとにExperimentを作成
    for exp in experiment_list:
        if data == exp.name:
            experiment_id = exp.experiment_id
            break
    else:
        experiment_id = client.create_experiment(data)
~略~

~略~
@hydra.main(config_path='config.yaml')# MLFlow単体のコードからの変更点
def workflow(cfg : omegaconf.DictConfig):# MLFlow単体のコードからの変更点
    ############## MLFlow単体のコードからの変更点###############

    data = cfg.data
    mecab_dic = cfg['mecab-dic']
    embedding = cfg['embedding']
    os.chdir(hydra.utils.get_original_cwd()) 
    client = mlflow.tracking.MlflowClient()
    #########################################################
    
    experiment_list = client.list_experiments()

    # データごとにExperimentを作成
    for exp in experiment_list:
        if data == exp.name:
            experiment_id = exp.experiment_id
            break
    else:
        experiment_id = client.create_experiment(data)


~略~

また、Hydraの特徴の1つに、Config Groupという仕組みでパラメータに階層性を持たせることが出来るということが挙げられます。例えば、以下のように前処理とモデルのパラメータを分けてYAMLの記述を行うことが出来ます。

preprop: 
     mecab-dic: ipadic
train:
    dropout1: 0.5
    lr: 0.1

コマンドライン上でパラメータを指定する場合は以下のように指定します。

main.py preprop.mecab-dic=ipadic train.lr=0.01

Hydraのハイパーパラメータサーチはグリッドサーチのみ対応しています。ベイズ最適化など高度な手法でハイパーパラメータサーチをやりたい場合はOptunaなど別のツールを使うか、組み合わせて使う必要があります。

おわりに

MLFlowとKedro、Hydraの簡単な紹介を行いました。最近のMLOps系ツールの発展は目覚ましく、注視していきたいところです。今後も、新たな知見が得られれば発信していきたいと思います。