MLflow:Basic

lastmod 2026-03-04created 2026-03-043 edits

Quickstart

MLflow Tracking Quickstart | MLflow

with Docker

Official MLflow Docker image | MLflow

How to apply MLflow

1. Installation

pip install mlflow

2. Changes to train.py

You can keep or replace the existing TensorBoard code. Key changes:

import mlflow
import mlflow.pytorch

Start the MLflow run in init.py

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,
})

Inside the train() loop

Add MLflow metrics alongside the existing TensorBoard self._writer.add_scalar calls:

# 기존 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)

At the end of train()

Save the model artifact after self._writer.close():

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

At the end of eval()

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

At the end of 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. Run the MLflow UI

mlflow ui --host 0.0.0.0 --port 5000

Open http://localhost:5000 in a browser to compare experiment results, search hyperparameters, and manage model artifacts.

Summary comparison

Item	TensorBoard (current)	MLflow (added)
Metric visualization	O	O
Hyperparameter comparison	Limited	O (table format)
Model versioning	X	O (artifacts)
Experiment search/filter	X	O
Model registry	X	O

Favorite site

author Zer0 sitemap