MLflow:Basic
Quickstart
with Docker
MLflow 적용 방법
1. 설치
2. train.py 수정 포인트
기존 TensorBoard 코드를 유지하거나 교체할 수 있다. 핵심 변경:
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() 종료 시
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 실행
브라우저에서 http://localhost:5000으로 접속하면 실험 결과 비교, 하이퍼파라미터 검색, 모델 아티팩트 관리가 가능하다.
요약 비교
| 항목 | TensorBoard (현재) | MLflow (추가) |
| 메트릭 시각화 | O | O |
| 하이퍼파라미터 비교 | 제한적 | O (표 형태) |
| 모델 버전 관리 | X | O (아티팩트) |
| 실험 검색/필터 | X | O |
| 모델 레지스트리 | X | O |