はじめに

Detectron2は物体検出やセグメンテーションをする上で色々と便利な機能が簡単に実装できるライブラリです。
ただ、使ってみるとモデルを構築する上で提供されていないものがチラホラあるので、掘り下げて実装してみました。
今回は実験管理ツールであるMLFlowで評価指標のログを取れるようにしてみます。
Detectron2はtensorboardをデフォルトでサポートはしていますがMLFlowを使うにはカスタムで実装しないといけません。
ただ、割と簡単に実装ができます。

環境

Python 3.10.13
CUDA 11.7
pytorch==1.13.1
detectron2==0.6

実装

Hookとして以下のようなクラスを定義します。

import mlflow
import torch
from detectron2.config import CfgNode
from detectron2.engine import HookBase


class MLflowTracker(HookBase):
    def __init__(self, cfg: CfgNode) -> None:
        super().__init__()
        self.cfg = cfg.clone()

    def before_train(self) -> None:
        with torch.no_grad():
            mlflow.set_tracking_uri(self.cfg.MLFLOW.TRACKING_URI)
            mlflow.set_experiment(self.cfg.MLFLOW.EXPERIMENT)
            mlflow.start_run(run_name=self.cfg.MLFLOW.RUN_NAME)
            for k, v in self.cfg.items():
                mlflow.log_param(k, v)

    def after_step(self) -> None:
        with torch.no_grad():
            latest_metrics = self.trainer.storage.latest()
            for k, v in latest_metrics.items():
                mlflow.log_metric(key=k, value=v[0], step=v[1])

    def log_artifacts(self, file_dir: str) -> None:
        mlflow.log_artifacts(file_dir)

    def log_artifact(self, artifact_path: str) -> None:
        mlflow.log_artifact(artifact_path)

定義したHookをTrainerのregister_hooksに渡します。

main.py

# mlflow
mlflow_tracker = MLflowTracker(config)

# train
trainer = Trainer(config)
trainer.register_hooks(hooks=[mlflow_tracker])

...

これで学習を行うとMLFlowで評価指標のログを取ることができます。

実験結果の確認

mlflow uiコマンドを実行して、実験結果を表示する画面を起動します。
デフォルトではhttp://127.0.0.1:5000で表示できます。

Artifactとしてファイルを保存

log_artifacts()で指定のディレクトリのファイルすべてを、log_artifactで指定のファイルをArtifactとして保存することができます。

main.py

# store artifacts in mlflow
mlflow_tracker.log_artifacts(config.OUTPUT_DIR)
mlflow_tracker.log_artifact("config/parameter_config.yaml")