MLflow Model API

머신러닝 모델을 개발하고 배포할 때 가장 중요한 것은 일관된 관리와 재사용성입니다.
MLflow의 Model API는 이러한 문제를 해결하기 위해 표준화된 모델 저장, 로드, 배포 방식을 제공합니다.
이번 포스팅에서는 MLflow Model API의 개념과 실제 코드 예제를 통해 그 원리를 살펴보겠습니다.

 

 

1. MLflow Model API란?

MLflow Model API는

머신러닝 모델을 다양한 프레임워크에서 일관된 방식으로 저장하고 불러오며, 배포할 수 있도록 표준화된 인터페이스를 제공하는 API입니다.

 

즉, 모델을 한 번 정의하면
PyFunc(Python Function Model), REST API, 또는 클라우드 환경 등 어디서든 호환 가능하게 재사용할 수 있습니다.

 

 

2. 주요 함수

MLflow Model API는 세 가지 핵심 함수를 중심으로 작동합니다.

 

1️⃣ save_model()

• 모델을 지정된 디렉토리에 저장합니다.
• 예를 들어, 학습이 끝난 모델을 로컬 폴더에 저장해두고 나중에 재사용할 때 활용됩니다.

2️⃣ log_model()

• 모델을 MLflow Tracking Server에 저장합니다.
• 이를 통해 여러 실험(run)을 관리하고, 모델 버전을 추적할 수 있습니다.

3️⃣ load_model()

• 저장된 모델을 로컬 디렉토리나 MLflow Registry에서 불러옵니다.
• 배포나 추론 시점에서 모델을 쉽게 복원할 수 있습니다.

 

3. 예제 코드: Scikit-learn 모델 관리

아래 예제는 scikit-learn의 Logistic Regression 모델을 학습하고,
MLflow API를 사용해 로컬 저장 및 트래킹 서버에 등록하는 과정을 보여줍니다.

import mlflow
import mlflow.sklearn
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris

# Load data
X, y = load_iris(return_X_y=True)

# Train a model
model = LogisticRegression()
model.fit(X, y)

# Save model locally
mlflow.sklearn.save_model(model, path="model_api")

# Log model to MLflow Tracking Server
with mlflow.start_run():
    mlflow.sklearn.log_model(model, artifact_path="artifact_model_api")

코드 설명

1. 데이터 로드:
   load_iris()를 통해 Iris 데이터셋을 불러옵니다.
2. 모델 학습:
   LogisticRegression() 모델을 훈련합니다.
3. 로컬 저장:
   save_model()로 모델을 "model_api" 폴더에 저장합니다.
4. MLflow 서버 로깅:
   start_run() 컨텍스트 안에서 log_model()을 호출하면
   MLflow Tracking Server에 모델 아티팩트가 저장됩니다.

 

4. 왜 MLflow Model API를 써야 할까?

✅ 프레임워크 독립성:
Scikit-learn, TensorFlow, PyTorch 등 다양한 프레임워크를 동일한 방식으로 다룸

✅ 모델 버전 관리:
각 실험별로 모델을 로깅하고, MLflow UI에서 버전을 비교 가능

✅ 손쉬운 배포:
PyFunc 형식으로 저장된 모델은 REST API나 Docker 컨테이너로 바로 배포 가능

✅ 팀 협업에 유리:
Tracking Server를 통해 여러 구성원이 실험 결과를 공유하고 재현 가능

 

 

정리하면

• save_model() → 로컬 저장
• log_model() → MLflow 서버 로깅
• load_model() → 모델 불러오기

 

MLflow Model API는 모델 저장과 배포를 단순화하면서도 일관성 있게 관리할 수 있는 강력한 도구입니다.
단 몇 줄의 코드로 모델을 저장하고, 추후 재사용이나 배포 시 손쉽게 불러올 수 있습니다.

다음 포스팅에서는 MLflow Model Registry를 통해
저장된 모델을 버전 관리하고, 프로덕션 환경으로 배포하는 방법을 다뤄보겠습니다.