mlflow

J. Hwang·2024년 12월 19일

AI 서비스를 개발할 때, 데이터 수집/라벨링/정제 작업과 다양한 모델 및 파라미터로 실험을 진행하고 가장 성능이 좋은 모델을 배포하게 된다. 이러한 일련의 과정은 항상 꼼꼼히 기록되고 관리되어야 한다.

모델을 관리한다고 함은, 1) 메타 데이터 (언제 만들어졌고, 어떤 데이터를 사용했는지 등), 2) 아티팩트 (모델의 학습된 결과물 (pickle, joblib)), 3) feature / data 를 기록하고 관리한다는 것이다.

mlflow는 머신러닝 라이프사이클 관리 플랫폼으로, 머신러닝 모델을 효율적으로 실험, 배포, 관리할 수 있도록 도와주는 툴이다. mlflow 사용법을 배워 모델 실험/개발 부터 배포까지 잘 관리해보자.

핵심 기능

실험 실행 / 관리 / 기록
registry : 실험한 모델을 저장하고 버전 관리. 다른 사용자들에게 공유하고 활용 가능하다.
serving : registry에 등록한 모델을 REST API 서버로 배포

설치

$ pip install mlflow==2.10.0

UI 실행

$ mlflow server --host 127.0.0.1 --port 8080
위 명령어를 실행하고 브라우저에서 localhost:8080으로 접속하면 mlflow 웹 페이지를 볼 수 있다.

experiment 생성

experiment는 하나의 프로젝트를 의미한다. (ex) hand bone image segmentation)
experiment는 여러 run (실행) 으로 구성된다.
experiment를 생성할 때에는 아래 명령어를 이용하면 된다.
$ mlflow experiments create --experiment-name hand_bone_segmentation
experiment가 성공적으로 생성되면 mlruns라는 폴더가 생성되고, 웹 페이지를 확인했을 때에도 experiment가 생성된 것을 확인할 수 있다.
experiment의 목록을 확인하고 싶다면, $ mlflow experiments search를 이용하면 된다.

프로젝트 메타 정보 저장

experiment를 생성했다면, 이 프로젝트가 어떤 환경에서 어떻게 실행되는지 등 메타 정보를 저장해야 관리하기에 용이할 것이다. 따라서 이러한 정보를 담은 MLProject, python_env.yaml 파일을 작성한다.

MLProject

name: project1

python_env: python_env.yaml

entry_points:
  main:
    parameters:
      regularization: {type: float, default: 0.1}
    command: "python train.py"

python_env.yaml

python: "3.9.13"

# 패키지를 빌드할 때 필요한 dependency
build_dependencies:
  - pip
  - setuptools
  - wheel==0.37.1
  
# 프로젝트 실행할 때 필요한 dependency
dependencies:
  - mlflow==2.10.0
  - scikit-learn==1.4.0
  - pandas==2.2.0
  - numpy==1.26.3
  - matplotlib==3.8.2

실험 (run)

mlflow run experiment_directory --experiment-name hand_bone_segmentation
위 명령어를 이용하면 실험을 할 수 있다. experiment_directory에는 MLproject 파일과 train.py 등 실험 (학습) 코드가 담긴 디렉토리 경로를 입력하면 된다.
-P 옵션으로 파라미터를 지정할 수 있다. ex) -P regularization=0.01 (단, 파라미터 명은 MLProject에 등록한 것을 사용해야 한다.)

run을 하면 웹 서버에 run 기록이 생기고 아래와 같은 정보들이 기록된다.

source : 실행한 프로젝트의 름
version : 실행 hash
start & end time
parameters
metrics
tags
artifacts : 실행 과정에서 생기는 여러 파일들 (이미지, 모델 pickle)

train.py 내에 어떤 항목을 기록할지 지정할 수도 있다. autolog를 활용하면 많은 항목들에 대해 자동으로 기록되는 것을 확인할 수 있다. 단, autolog는 모든 프레임워크에서 사용이 가능한 것은 아니라는 점에 유의해야 한다. (autolog의 자세한 것은 여기 참조)

import mlflow
import mlflow.sklearn

# 개별 항목에 대해 지정
mlflow.log_param("l1_ratio":0.1)

# autolog
mlflow.sklearn.autolog()

with mlflow.start_run() as run:
	model.fit(x, y)

모델 다운로드

아래 코드를 이용하면 run id를 입력해 모델 아티팩트를 다운로드할 수 있다.

from mlflow import artifacts

def download_model(run_id, model_name="model"):
	artifact_uri=f"runs:/{run_id}/{model_name}"
    artifacts.download_artifacts(artifact_uri, dst_path=".")