모델 학습에 MLflow 적용하기

 

 

MLflow란

MLflow는 데이터 과학 및 머신러닝 프로젝트를 위한 오픈소스 플랫폼이다. MLflow는 실험 관리, 모델 등록, 모델 서빙, 그리고 모델 배포 등 머신러닝 생명주기의 여러 단계를 관리할 수 있는 기능을 제공한다. 이 플랫폼은 Databricks에 의해 개발되었고, Python, R, 그리고 Java 같은 다양한 프로그래밍 언어를 지원한다.

 

 

MLflow의 주요 기능

1. MLflow Tracking: 실험을 추적하고 파라미터, 코드 버전, 메트릭, 그리고 결과를 기록한다.
2. MLflow Projects: 프로젝트 환경을 코드로 정의하여 재현성을 높인다.
3. MLflow Models: 다양한 ML 라이브러리에서 사용되는 모델을 포장(packaging)하고 배포하기 위한 표준 포맷을 제공한다.
4. MLflow Registry: 모델을 중앙 집중식으로 저장하고, 모델 버전 관리를 수행한다.

 

 

MLflow 설치

pip install mlflow

 

Code

해당 예시에서는 MNIST 데이터셋을 사용하여 간단한 CNN 모델을 학습하고 MLflow를 적용해 본다.

 

필요 라이브러리

import mlflow
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

 

MNIST 데이터셋 로드

transform = transforms.Compose([transforms.ToTensor(), 
                                transforms.Normalize((0.1307,), (0.3081,))]
                               )
train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('data', train=False, download=True, transform=transform)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1000, shuffle=False)

 

Accuracy 계산 함수

def calculate_accuracy(model, data_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in data_loader:
            output = model(data)
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    return correct / total

 

하이퍼파라미터 설정

learning_rate = 0.01
epochs = 5

 

MLflow 실험 시작

mlflow.set_tracking_uri("http://127.0.0.1:5000")
mlflow.set_experiment("MNIST_Pytorch")

with mlflow.start_run():
    model = Net()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)

    # 하이퍼파라미터 기록
    mlflow.log_param("learning_rate", learning_rate)
    mlflow.log_param("epochs", epochs)

    # 학습 루프
    for epoch in range(epochs):
        model.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            if batch_idx % 100 == 0:
                print(f"Epoch {epoch + 1}/{epochs}, Batch {batch_idx}/{len(train_loader)}, Loss: {loss.item():.4f}")

        train_accuracy = calculate_accuracy(model, train_loader)
        test_accuracy = calculate_accuracy(model, test_loader)               

        # metric 기록
        mlflow.log_metric("train_loss", loss.item(), step=epoch)
        mlflow.log_metric("train_accuracy", train_accuracy, step=epoch)
        mlflow.log_metric("test_accuracy", test_accuracy, step=epoch)

        print(f"Epoch {epoch + 1}/{epochs}, Train Accuracy: {train_accuracy:.4f}, Test Accuracy: {test_accuracy:.4f}")

    # 모델 저장
    mlflow.pytorch.log_model(model, "mnist_model")

 

결과

terminal에서 mlflow ui를 입력하고 다른 terminal에서는 모델을 학습하면 된다.

위에 설정한 127.0.0.1:5000에 접속하며 아래와 같은 화면을 볼 수 있다.

 

 

해당 모델을 눌러보면 hyperparameter를 어떻게 설정했는지와 각 epoch마다 accuracy, loss 등을 볼 수 있다.

 

기존에 모델을 학습할 때, hyperparameter와 accuracy를 모델 파일명에 함께 저장을 했다. 아직 MLflow의 모든 기능을 다 활용하지는 못하지만 MLflow를 사용하면 실험에 대한 메타데이터들을 관리할 수 있어 더욱 효율적이고 체계적으로 실험 관리가 이루어질 수 있을 거 같다.

 

참고자료

https://mlflow.org/

MLflow | MLflow