MLflow 사용해보기

한승수·2025년 12월 17일

AI 코딩 팁

목록 보기

12/13

머신러닝 프로젝트를 진행하다 보면 "저번 실험에서 학습률(Learning Rate)을 얼마로 줬더라?", "이 모델 가중치 파일은 어떤 데이터셋으로 학습한 거지?" 하는 혼란에 빠지기 쉽습니다. 이런 실험 관리의 지옥에서 우리를 구해줄 도구가 바로 MLflow입니다.

MLflow의 개념부터 실제 적용법까지 핵심만 정리해 드립니다.

1. MLflow란 무엇인가?

MLflow는 머신러닝 라이프사이클 전체를 관리하기 위한 오픈소스 플랫폼입니다. 모델 학습 과정의 모든 기록을 체계적으로 저장하고, 나중에 이를 쉽게 비교하거나 배포할 수 있도록 돕습니다.

2. 어떨 때 쓰는 게 좋은가?

하이퍼파라미터 튜닝 시: 수십 번의 실험 결과(Loss, Accuracy 등)를 엑셀에 수동으로 적는 대신 자동으로 기록하고 싶을 때.
팀 프로젝트 수행 시: 팀원들이 각자 돌린 실험 결과를 한곳에 모아 비교하고 싶을 때.
모델 재현성이 중요할 때: 특정 결과가 나왔던 당시의 코드 버전, 데이터셋, 파라미터를 그대로 복원해야 할 때.
모델 배포 시: 학습된 모델을 표준화된 방식으로 저장하여 서빙 환경에 쉽게 올리고 싶을 때.

3. 어떻게 쓰는 것인가? (일반적인 PyTorch 예시)

PyTorch에서 MLflow를 사용하는 가장 표준적인 방법은 mlflow.start_run() 컨텍스트 매니저를 사용하는 것입니다.


import mlflow
import mlflow.pytorch
import torch

# 1. 서버 연결 및 실험 설정
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("My_First_Experiment")

with mlflow.start_run():
    # 2. 파라미터 기록 (설정값)
    mlflow.log_param("lr", 0.01)
    mlflow.log_param("batch_size", 32)

    # --- 모델 학습 루프 ---
    for epoch in range(10):
        loss = 0.5 / (epoch + 1) # 예시 loss
        # 3. 메트릭 기록 (변하는 지표)
        mlflow.log_metric("loss", loss, step=epoch)

    # 4. 모델 저장 (가중치 파일)
    mlflow.pytorch.log_model(model, "model")

4. Huggingface 모델을 쓰는 경우 (Trainer 활용)

Huggingface Trainer를 사용 중이라면 코드가 훨씬 단순해집니다. Trainer 내부에 이미 MLflow 콜백이 내장되어 있기 때문입니다.


from transformers import TrainingArguments, Trainer

# 별도의 로깅 코드 없이 설정만으로 가능
training_args = TrainingArguments(
    output_dir="./results",
    report_to="mlflow",      # 'mlflow'를 리포팅 대상으로 지정
    logging_steps=10,        # 10스텝마다 자동으로 MLflow에 기록
    run_name="my_hf_run"     # MLflow에 표시될 실행 이름
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    # ...
)

# 학습 시작 시 자동으로 MLflow Run이 생성되고 기록됨
trainer.train()