MLflow Tracking

MLflow Tracking

머신러닝 코드를 실행할 때 파라미터, 코드 버전, 메트릭, 출력 파일을 기록하고 나중에 결과를 시각화하는 데 쓰는 API와 UI가 MLflow Tracking이에요. Python, REST, R, Java API를 모두 제공해요. 처음 쓴다면 quickstart 튜토리얼을 먼저 진행해 보는 걸 권장해요.

출처: MLflow Tracking

실행(Run)과 실험(Experiment)

MLflow Tracking은 run(실행)이라는 개념으로 구성돼요. run은 데이터 사이언스 코드의 한 실행, 예를 들어 python train.py 한 번 실행이에요. 각 run은 메타데이터(메트릭·파라미터·시작·종료 시간 등)와 아티팩트(모델 가중치·이미지 같은 출력 파일)를 기록해요.

모델(run 중 생성된 학습된 머신러닝 아티팩트)은 run과 비슷하게 자체 메타데이터와 아티팩트를 갖고 있어요. 그리고 실험(Experiment)은 특정 작업의 run과 모델을 묶는 그룹이에요. CLI·API·UI로 실험을 만들 수 있어요.

Tracking API

mlflow.start_run()으로 새 run을 시작하고, mlflow.log_param()·mlflow.log_metric() 같은 로깅 함수로 파라미터와 메트릭을 기록해요.

import mlflow
with mlflow.start_run():
    mlflow.log_param("lr", 0.001)
    # Your ml code...
    mlflow.log_metric("val_loss", val_loss)

Autologging

Autologging은 명시적 log 문 없이 메트릭·파라미터·모델을 로깅하게 해 주는 초고속 설정이에요. 학습 코드 앞에 mlflow.autolog()을 호출하기만 하면 돼요. Scikit-learn, XGBoost, PyTorch, Keras, Spark 등 인기 라이브러리를 지원해요.

import mlflow
mlflow.autolog()
# Your training code...

추가 설정 없이 MLflow Tracking은 데이터를 로컬 mlruns 디렉토리에 로깅해요. 팀과 결과를 공유하려면 원격 DB·클라우드 스토리지로 설정을 바꾸면 돼요.

모델 검색

MLflow 3는 mlflow.search_logged_models()로 강력한 모델 검색 기능을 제공해요. SQL 같은 문법으로 성능 메트릭·파라미터·모델 속성에 기반해 실험 전반에서 특정 모델을 찾을 수 있어요.

import mlflow
# Find high-performing models across experiments
top_models = mlflow.search_logged_models(
    experiment_ids=["1", "2"],
    filter_string="metrics.accuracy > 0.95 AND params.model_type = 'RandomForest'",
    order_by=[{"field_name": "metrics.f1_score", "ascending": False}],
    max_results=5,
)
# Get the best model for deployment
best_model = mlflow.search_logged_models(
    experiment_ids=["1"],
    filter_string="metrics.accuracy > 0.9",
    max_results=1,
    order_by=[{"field_name": "metrics.accuracy", "ascending": False}],
    output_format="list",
)[0]
# Load the best model directly
loaded_model = mlflow.pyfunc.load_model(f"models:/{best_model.model_id}")

핵심 특징을 정리하면:

  • SQL 같은 필터링: metrics., params. 접두사와 속성으로 복잡한 쿼리 구성
  • 데이터셋 인지 검색: 공정한 모델 비교를 위해 특정 데이터셋 기반으로 메트릭 필터링
  • 유연한 정렬: 여러 기준으로 정렬해 최고 모델 찾기
  • 직접 모델 로딩: models:/<model_id> URI 형식으로 즉시 모델 접근

MlflowClient

Tracking UI의 모든 기능을 프로그래밍 방식으로 쓰려면 MlflowClient를 사용해요. 예를 들어 실험에서 검증 손실이 가장 낮은 run을 검색할 수 있어요.

client = mlflow.tracking.MlflowClient()
experiment_id = "0"
best_run = client.search_runs(experiment_id, order_by=["metrics.val_loss ASC"], max_results=1)[0]
print(best_run.info)  # {'run_id': '...', 'metrics': {'val_loss': 0.123}, ...}

모델 체크포인트·메트릭 추적

MLflow 3는 단일 run 안에서 여러 모델 체크포인트를 로깅하고 서로 다른 데이터셋에 대한 성능을 추적하게 해 줘요. 딥러닝에서 학습 단계별 체크포인트를 저장·비교할 때 특히 유용해요. 모델 로깅 함수의 step 파라미터로 학습 중 여러 지점의 체크포인트를 로깅하고, 각 로그된 모델은 나중에 참조할 수 있는 고유 model ID를 받아요.

import mlflow
import mlflow.pytorch
with mlflow.start_run() as run:
    for epoch in range(100):
        # Train your model
        train_model(model, epoch)
        # Log model checkpoint every 10 epochs
        if epoch % 10 == 0:
            model_info = mlflow.pytorch.log_model(
                pytorch_model=model,
                name=f"checkpoint-epoch-{epoch}",
                step=epoch,
                input_example=sample_input,
            )
        # Log metrics linked to this specific model checkpoint
        accuracy = evaluate_model(model, validation_data)
        mlflow.log_metric(key="accuracy", value=accuracy, step=epoch, model_id=model_info.model_id)

MLflow 3는 메트릭을 특정 모델 체크포인트·데이터셋에 연결해 모델 성능 추적성을 높여 줘요.

# Create a dataset reference
train_dataset = mlflow.data.from_pandas(train_df, name="training_data")
# Log metric with model and dataset links
mlflow.log_metric(key="f1_score", value=0.95, step=epoch, model_id=model_info.model_id, dataset=train_dataset)

MLflow 3는 run ID 대신 model ID를 쓰는 새 모델 URI 형식을 도입했어요. 더 직접적으로 모델을 참조하죠.

# New MLflow 3 model URI format
model_uri = f"models:/{model_info.model_id}"
loaded_model = mlflow.pyfunc.load_model(model_uri)
# This replaces the older run-based URI format:
# model_uri = f"runs:/{run_id}/model_path"

이 방식의 장점은 이래요.

  • 직접 모델 참조: run ID와 아티팩트 경로를 몰라도 돼요.
  • 더 나은 모델 수명주기 관리: 각 체크포인트가 고유 식별자를 가져요.
  • 개선된 모델 비교: 같은 run 안 체크포인트 비교가 쉬워요.
  • 강화된 추적성: 모델·메트릭·데이터셋 사이의 명확한 연결.

데이터셋 추적·UI

MBflow는 모델 학습 이벤트와 연관된 데이터셋을 추적할 수 있어요. mlflow.log_input() API로 이 메타데이터를 저장해요. Tracking UI에서는 실험·run·모델을 시각적으로 탐색하고, 파라미터·메트릭 값으로 run을 검색하고, run 메트릭을 시각화하고, run 결과(아티팩트·메타데이터)를 다운로드할 수 있어요.

로컬 mlruns 디렉토리에 로깅했다면 그 상위 디렉토리에서 다음을 실행하고 브라우저로 http://127.0.0.1:5000에 접속해요.

mlflow server --port 5000

더 알아보기