MLflow:Basic

Quickstart

MLflow Tracking Quickstart | MLflow

with Docker

Official MLflow Docker image | MLflow

MLflow 적용 방법

1. 설치

pip install mlflow

2. train.py 수정 포인트

기존 TensorBoard 코드를 유지하거나 교체할 수 있다. 핵심 변경:

import mlflow
import mlflow.pytorch

init.py 에서 MLflow run 시작

mlflow.set_experiment("vrs-vehicle-recommender")
mlflow.start_run()

# 하이퍼파라미터 기록
mlflow.log_params({
    "model_version": self._model_version,
    "test_size": self._test_size,
    "random_state": self._random_state,
    "batch_size": self._batch_size,
    "lr": self._lr,
    "epochs": self._epochs,
    "shuffle": self._shuffle,
})

train() 루프 안에서

기존 TensorBoard 의 self._writer.add_scalar 부분에 MLflow 메트릭 추가:

# 기존 TensorBoard 유지 가능
self._writer.add_scalar("Train/Loss", train_loss, epoch)
self._writer.add_scalar("Train/Accuracy", train_accuracy, epoch)

# MLflow 메트릭 추가
mlflow.log_metrics({
    "train_loss": train_loss,
    "train_accuracy": train_accuracy,
}, step=epoch)

train() 종료 시

self._writer.close() 이후에 모델 아티팩트 저장:

# 모델 아티팩트 저장
mlflow.pytorch.log_model(self._model, "model")
mlflow.log_metric("final_train_loss", train_loss)

eval() 종료 시

mlflow.log_metrics({
    "test_loss": test_loss,
    "test_accuracy": test_accuracy,
})

run() 종료 시

def run(self) -> None:
    self.init()

    if not self._only_eval:
        self.train()

    if not self._without_eval:
        self.eval()

    mlflow.end_run()

3. MLflow UI 실행

mlflow ui --host 0.0.0.0 --port 5000

브라우저에서 http://localhost:5000으로 접속하면 실험 결과 비교, 하이퍼파라미터 검색, 모델 아티팩트 관리가 가능하다.

요약 비교

항목	TensorBoard (현재)	MLflow (추가)
메트릭 시각화	O	O
하이퍼파라미터 비교	제한적	O (표 형태)
모델 버전 관리	X	O (아티팩트)
실험 검색/필터	X	O
모델 레지스트리	X	O