MLflow Model Customization — 모델 커스터마이징으로 유연한 확장성 확보하기

MLflow는 단순히 모델을 저장하고 배포하는 도구가 아닙니다.
MLflow의 진짜 강점은 기존 머신러닝 프레임워크를 넘어서 모델 로직을 확장(customize) 할 수 있다는 점입니다.
이번 포스팅에서는 MLflow Model Customization 개념과, 사용자 정의 모델을 만드는 방법을 예제 코드와 함께 살펴보겠습니다.

 

 

1. MLflow Model Customization이란?

MLflow는 사용자가 모델의 동작 방식을 자유롭게 확장할 수 있도록 커스터마이징(customizing) 기능을 제공합니다.

즉, 단순히 학습 모델만 저장하는 것이 아니라,
전처리(Preprocessing), 후처리(Post-processing), 앙상블 로직(Ensemble),
**사용자 정의 추론(Inference Logic)**까지 모델 내부에 통합할 수 있습니다.

 

이러한 기능을 활용하면, 복잡한 머신러닝 파이프라인도 하나의 모델로 캡슐화할 수 있습니다.

 

2. 왜 MLflow 모델을 커스터마이징할까?

MLflow 커스터마이징은 “단순한 모델 저장”을 넘어, 데이터 파이프라인과 로직 전체를 하나의 모델 단위로 관리하기 위해 사용됩니다.

 

1️⃣ Add Preprocessing Logic

모델 입력 전에 데이터를 가공할 수 있습니다.
예: 데이터 정규화, 피처 엔지니어링, 범주형 변수 인코딩 등

 

2️⃣ Modify Post-processing

모델의 예측 결과를 후처리하여 사람이 이해하기 쉬운 형태로 변환합니다.
예: 예측 확률 → 범주 레이블로 변환

 

3️⃣ Combine Multiple Models

여러 모델을 하나의 파이프라인으로 결합할 수 있습니다.
예: 앙상블 모델(Ensemble Model) 형태로 구현 가능

 

4️⃣ Extend Beyond MLflow Flavors

기존 MLflow 내장 flavor(예: mlflow.sklearn, mlflow.pytorch)로 처리할 수 없는
딥러닝 프레임워크나 비즈니스 로직을 직접 정의할 수 있습니다.

 

 

3. mlflow.pyfunc.PythonModel — 사용자 정의 모델의 핵심

MLflow에서는 mlflow.pyfunc.PythonModel 클래스를 상속받아
**사용자 정의 모델(Custom MLflow Model)**을 만들 수 있습니다.

이 클래스를 활용하면,

• 모델 내부의 전처리 로직,
• 예측(Predict) 과정,
• 필요 시 후처리(Post-processing) 단계까지 직접 정의할 수 있습니다.

 

4. 예제 

1️⃣  전처리 포함 사용자 정의 모델

# Define a custom model with preprocessing
class CustomModel(PythonModel):
    def __init__(self, sklearn_model):
        self.sklearn_model = sklearn_model
        self.scaler = StandardScaler()  # Example of a preprocessing step

    def fit(self, X_train, y_train):
        X_scaled = self.scaler.fit_transform(X_train)
        self.sklearn_model.fit(X_scaled, y_train)

    def predict(self, context, model_input):
        # Apply preprocessing
        model_input_scaled = self.scaler.transform(model_input)
        return self.sklearn_model.predict(model_input_scaled)

코드 설명

• __init__에서 Scaler(정규화기)와 sklearn 모델을 초기화합니다.
• ✅  PythonModel을 상속하는 이유는 MLflow가 해당 객체를 예측 가능한 모델로 인식하고 관리할 수 있게 하기 위해서입니다. 
• fit()에서 입력 데이터를 스케일링 후 모델을 학습시킵니다.
• predict()에서는 추론 시 동일한 전처리를 적용합니다.

즉, 전처리 + 모델 학습 + 추론 로직이 하나의 클래스 안에 통합된 형태입니다.

 

 2️⃣  예제 2: 커스텀 모델 저장하기

# Train model
model = LogisticRegression(max_iter=200)
custom_model = CustomModel(model)
custom_model.fit(X, y)

# Save as an MLflow model
mlflow.pyfunc.save_model(
    path="custom_mlflow_model",
    python_model=custom_model
)

이 코드는 CustomModel을 MLflow 포맷으로 저장하는 과정입니다.
save_model() 호출 시, MLflow는 내부적으로 다음 파일들을 함께 생성합니다:

• MLmodel (메타데이터)
• python_model.pkl (직렬화된 사용자 모델)
• conda.yaml (환경 정보)

이제 이 모델은 MLflow Tracking Server나 REST API를 통해 쉽게 로드(load_model)하여 사용할 수 있습니다.

 

5. 커스터마이징을 통한 확장의 장점

항목	설명
유연한 전처리 포함	데이터 정규화, 인코딩 등을 모델 내부에 포함
일관된 추론 파이프라인	예측 시에도 동일한 로직 적용
엔드투엔드 파이프라인 구성	모델 + 전처리 + 후처리 통합 관리
엔터프라이즈 수준 확장성	내부 비즈니스 로직까지 포함 가능

 

 

정리하면

 

MLflow Model Customization은 단순히 모델을 저장하는 기능을 넘어,
데이터 전처리부터 추론까지 전체 파이프라인을 하나의 모델 객체로 관리할 수 있게 해줍니다.

✅ 핵심 요약

• mlflow.pyfunc.PythonModel을 상속받아 커스텀 모델 정의
• 전처리, 후처리, 앙상블 등 복잡한 로직 통합 가능
• mlflow.pyfunc.save_model()로 저장 후 MLflow 환경에서 바로 재사용 가능