[MLOps]MLFlowとFastLabelを連携し、実験管理をする

はじめに

AIプロジェクトは、データの前処理からモデルのトレーニング、評価、デプロイメントまで多くの工程が含まれています。プロジェクトが成長するにつれて、実験、データやモデルのバージョン管理、パラメータの調整、結果の追跡など、管理上の課題が増えてきます。このような課題に対処するために、MLflowとFastLabelを連携させてMLOpsを回し、モデル開発の実験管理をします。

アーキテクチャイメージ

S3にデータを蓄積し、FastLabel蓄積した画像の可視化とアノテーションの実施をします。
その後、学習を回し実験管理を以下の役割で行います。

• 画像やアノテーションデータの調整・バージョン管理　→　FastLabel
• ハイパーパラメータの調整・記録　→　MLflow

アーキテクチャの全体像として一般的に下画像のような形が想定されます。
こちらの図ではAWSを基盤としていますが、ハンズオンではColabで完結できるようになっています。

FastLabelにおけるデータの管理について

MLOpsにおいて、アノテーションや学習データの管理は非常に重要です。適切な学習データの管理は、モデルのトレーニングと評価における再現性、品質の維持、セキュリティを確保するために必要です。
FastLabelでは、以下の機能を提供しMLOpsを回す上で必要なデータの管理をすることができます。

• データの収集・蓄積
  簡単にFastLabelにデータを上げることができ、自動化することができます。FastLabelに収集したデータや統計を確認することができます。
• データの前処理
  画像へのオーグメンテーションや、リサイズをすることができます。
• データバージョニング
  学習データのバージョン管理を行います。また、アノテーションデータの品質を監視し、品質の低下を検出した場合に対処します。アノテーション品質に対する品質を管理できます。
• データの分割
  データをトレーニング、検証、テストセットに分割します。この分割は一貫性を保つためにバージョン管理されるべきです。
• データセットのドキュメンテーション
  データセットに関する情報（スキーマ、特徴量の説明、メタデータ）をドキュメント化します。これはチーム内のコラボレーションとプロジェクトの透明性を維持することに役立ちます。

FastLabelを利用することで、MLOpsにおける画像や動画の管理が容易にでき、モデルの分析やバージョン管理を実現することが可能できます。FastLabelについてのドキュメントはこちら

MLOpsの実験管理について

MLOps（Machine Learning Operations）は、機械学習プロジェクトを効率的かつ持続可能に展開、管理、および運用するためのプロセスです。本番環境へのデプロイメント、モデルの監視、メンテナンス、更新など、機械学習プロジェクト全体のライフサイクルをカバーします。

MLOpsにおける実験管理は、機械学習プロジェクト内で実行される実験やトレーニングの管理、トラッキング、および文書化を指します。実験管理のツール群は、プロジェクトの再現性、バージョン管理、モデルのトラッキング、パフォーマンスの監視などを強化し、プロジェクト全体の透明性を高めます。以下は、MLOpsにおける実験管理の主要なOSSツールです。

MLflow

MLflowは、オープンソースの機械学習プロジェクト管理プラットフォームで、以下の3つの主要なコンポーネントから成り立っています。

• Tracking: MLflow Trackingは、実験の実行を記録および追跡し、実験のパラメータ、メトリクス、アーティファクト（モデルなど）を保存します。これにより、実験結果を簡単に比較し、バージョニングされた実験結果を管理できます。

• Projects: MLflow Projectsは、コード、データ、およびハイパーパラメータを一つのパッケージにまとめ、再現可能なモデルトレーニングプロセスを提供します。プロジェクトをパッケージ化し、簡単に共有できます。

• Models: MLflow Modelsは、モデルの追跡、バージョニング、デプロイメントをサポートします。モデルを一元管理し、さまざまなプラットフォームでデプロイできるようにします。

MLflowは、機械学習エンジニアとデータサイエンティストが実験管理とモデル管理を行うための便利なツールで、モデルの開発からデプロイメントまでのプロセスを強化します。

kubeflow

Kubeflowは、Kubernetes上で機械学習ワークロードをデプロイおよび管理するためのオープンソースプラットフォームです。Kubeflowは、機械学習モデルの開発、トレーニング、デプロイメント、監視、スケーリングを支援します。Kubeflowの主要な機能は、Pipelines、Katib、Servingなどがあります。
Kubeflowは、スケーラビリティ、可用性、セキュリティを重視する企業や組織にとって非常に価値のあるツールです。

ハンズオン

今回は物体検出のモデルを用いて、FastLabelとmlflowの連携をします。
実行環境は、GoogleのColabを使います。

必要なライブラリのインストールとドライブへの接続

fastlabel python SDKやmlflowをインストールします。

import subprocess
import time
from google.colab import drive
drive.mount('/content/drive')
!pip install ultralytics
!pip install mlflow --quiet
!pip install pyngrok
!pip install sagemaker
!pip install boto3
!pip install pydantic==1.10.9
!pip install fastlabel
     
import os
import shutil
import pprint
import fastlabel
from random import random, randint
import mlflow.sklearn
from mlflow import log_metric, log_param, log_artifacts, log_input
from sklearn.ensemble import RandomForestRegressor
from mlflow.tracking import MlflowClient
import warnings
warnings.filterwarnings("ignore")
print("mlflow version: ",mlflow.__version__)

環境変数の設定

FastLabelアクセストークン、ngrokのトークン、mlflowデータの保存先を指定します。
FastLabel Python SDKを用いるためのアクセストークンの取得方法はこちらをご確認下さい。
ngrokについては無料でご利用いただけます。

FASTLABEL_ACCESS_TOKEN=""
MLFLOW_BASE_PATH="/content/drive/MyDrive/mlflow"
NGROK_AUTH_TOKEN = ""

%cd "$MLFLOW_BASE_PATH"
cwd = os.getcwd()
print("current working directory:", cwd)
ngrok.kill()
ngrok.set_auth_token(NGROK_AUTH_TOKEN)
ngrok_tunnel = ngrok.connect(addr="5000", proto="http", bind_tls=True)
print("MLflow Tracking UI:", ngrok_tunnel.public_url)
local_registry = "sqlite:///mlruns.db"
os.environ.setdefault("MLFLOW_TRACKING_URI", local_registry)
mlflow.set_tracking_uri(local_registry)
get_ipython().system_raw("mlflow ui --backend-store-uri sqlite:///mlruns.db --port 5000 &")
os.environ.setdefault("FASTLABEL_ACCESS_TOKEN",FASTLABEL_ACCESS_TOKEN)

FastLabelのデータを連携

FastLabelのプロジェクトからデータを連携します。
trainに使うタスクに「train」タグ、
valに使うタスクに「val」タグを付ける必要があります。

import os
import time
import urllib.request
from multiprocessing import Pool, cpu_count
import fastlabel


DATA_DOWNLOAD_PATH="/content/drive/MyDrive/sample" 
WORKSPACE_SLUG="workspace-8591526036"
PROJECT_SLUG = "worker-bbox"

client = fastlabel.Client()
def download_image(args):
    task:dict=args[0]
    output_dir_path=args[1]
    urllib.request.urlretrieve(task["url"],  os.path.join(output_dir_path,task["name"]))


def get_yolo_tasks(project_slug:str,output_dir_path:str,tags:list[str]) -> list:
    # Iterate pages until new tasks are empty.
    all_tasks = []
    offset = None
    while True:
        time.sleep(1)
        tasks = client.get_image_tasks(
            project=project_slug, limit=1000, offset=offset,tags=tags
        )
        all_tasks.extend(tasks)
        if len(tasks) > 0:
            offset = len(all_tasks)  # Set the offset
        else:
            break
    client.export_yolo(project=project_slug, tasks=all_tasks, output_dir=output_dir_path)
    with Pool(cpu_count()) as p:
        p.map(download_image, list(map(lambda x:[x,os.path.join(output_dir_path,"annotations")],all_tasks)))
    return all_tasks


IMAGE_DIR = os.path.join(DATA_DOWNLOAD_PATH,"annotations")
TRAIN_YOLO_DIR_PATH=os.path.join(DATA_DOWNLOAD_PATH,"train")
VAL_YOLO_DIR_PATH=os.path.join(DATA_DOWNLOAD_PATH,"val")
get_yolo_tasks(project_slug=PROJECT_SLUG,output_dir_path=TRAIN_YOLO_DIR_PATH,tags=["train"])
get_yolo_tasks(project_slug=PROJECT_SLUG,output_dir_path=VAL_YOLO_DIR_PATH,tags=["val"])
annotations = list(map(lambda x:x["value"],client.get_annotations(project=PROJECT_SLUG)))
FASTLABEL_DATA_URI=f"https://app.fastlabel.ai/{WORKSPACE_SLUG}/{PROJECT_SLUG}"

学習パラメータの設定

学習率、バッチサイズ、エポック数、隠れ層のユニット数、正則化項の重みなどのハイパーパラメータを設定します。

yaml_data_file_path=os.path.join(DATA_DOWNLOAD_PATH,"data.yaml")

data = {
    'train':TRAIN_YOLO_DIR_PATH or os.path.join(DATA_DOWNLOAD_PATH,'train'),
    'val': VAL_YOLO_DIR_PATH or os.path.join(DATA_DOWNLOAD_PATH,'val'),
    'nc': len(annotations),
    'names': annotations
}
model_config={
    "model":"yolov8n.pt",
    "data":yaml_data_file_path,
    "epochs":1,
    "show_labels":True,
    "project":DATA_DOWNLOAD_PATH,
    "name":"model",
}

with open(yaml_data_file_path, 'w') as file:
    yaml.dump(data, file)

学習の実行

学習の実行をします。

MLFLOW_EXPERIMENT_NAME="fl-sample-experiment"

from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.utils.callbacks.mlflow import on_pretrain_routine_end, on_fit_epoch_end, on_train_end
from mlflow.data.filesystem_dataset_source import FileSystemDatasetSource
import mlflow.pytorch
import mlflow
import inspect
from ultralytics import YOLO
import yaml
from mlflow.data.http_dataset_source import HTTPDatasetSource

import json
from typing import Any, Dict, Optional

from mlflow.data.dataset_source import DatasetSource
from mlflow.entities import Dataset as DatasetEntity
from mlflow.utils.annotations import experimental

class FastLabelDataset(Dataset):
    def __init__(
        self, source: DatasetSource, name: Optional[str] = None, digest: Optional[str] = None
    ):
        self._name = name
        self._source = source
        self._digest = digest or self._compute_digest()
    def _compute_digest(self) -> str:
        return "FastLabel Dataset"
    def _to_dict(self, base_dict: Dict[str, str]) -> Dict[str, str]:
        return base_dict
    def to_json(self) -> str:
        base_dict = {
            "name": self.name,
            "digest": self.digest,
            "source": self._source.to_json(),
            "source_type": self._source._get_source_type(),
        }
        return json.dumps(self._to_dict(base_dict))
    def name(self) -> str:
        if self._name is not None:
            return self._name
        else:
            return "FastLabel Dataset"
    def digest(self) -> str:
        return self._digest
    def source(self) -> DatasetSource:
        return self._source
    def profile(self) -> Optional[Any]:
        return {}
    def schema(self) -> Optional[Any]:
        return {}
    def _to_mlflow_entity(self) -> DatasetEntity:
        dataset_json = json.loads(self.to_json())
        return DatasetEntity(
            name=dataset_json["name"],
            digest=dataset_json["digest"],
            source_type=dataset_json["source_type"],
            source=dataset_json["source"],
            schema=dataset_json.get("schema"),
            profile=dataset_json.get("profile"),
        )

os.environ.setdefault(
    "MLFLOW_EXPERIMENT_NAME", MLFLOW_EXPERIMENT_NAME
)
experiment=mlflow.get_experiment_by_name(MLFLOW_EXPERIMENT_NAME)
if experiment is None:
  experiment_id = mlflow.create_experiment(MLFLOW_EXPERIMENT_NAME)
  experiment=mlflow.get_experiment(experiment_id)

mlflow.set_experiment(experiment_id=experiment.experiment_id)

def get_function_file_path(function_name):
    function_obj = globals().get(function_name)
    if function_obj:
        return inspect.getsourcefile(function_obj)
    else:
        return None

with mlflow.start_run():
    mlflow.set_tag("fastlabel-project", FASTLABEL_DATA_URI)
    dataset=FastLabelDataset(HTTPDatasetSource(FASTLABEL_DATA_URI),name=FASTLABEL_DATA_URI)
    mlflow.log_input(dataset)

    # MLflow  Example
    # mlflow.log_param("my", "param")
    # mlflow.log_metric("score", 100)
    # yolov8 supports mlflow.

    trainer = DetectionTrainer(overrides=model_config)
    trainer.train()

MLflowデータの確認

MLflowでの実験結果を以下のように確認をすることができます。

DatasetやTagに、FastLabelのURLを登録することでパラメータとデータのチューニングを容易にすることができます。

おわりに

AIプロジェクトの成功には、データの適切な管理とモデルの実験管理が不可欠です。FastLabelとMLflow、Kubeflowの組み合わせは、この重要な課題に対処するための強力なツールセットを提供します。

FastLabelを使用することで、データの収集、前処理、品質管理、バージョン管理が効率的に行えます。また、MLflowを導入することで、モデルの実験管理、トレーニング、デプロイメント、監視など、MLOpsの各段階での作業を効率的に管理できます。

AIプロジェクトの規模が拡大するにつれて、データとモデルの管理は複雑化しますが、適切なツールとプラクティスを採用することで、プロジェクトの品質と効率を向上させることができます。FastLabel、MLflow、Kubeflowなどのツールは、MLOpsの実現に向けたリソースです。