PyTorch-Lightning을 사용하여 모델 가중치를 MLflow 추적 서버에 저장하는 방법

2024-07-27

PyTorch-Lightning을 사용하여 모델 가중치를 MLflow 추적 서버에 저장하는 방법

필수 조건:

Python 3.6 이상
PyTorch 1.7 이상
PyTorch-Lightning 1.0 이상
MLflow 1.0 이상

단계:

MLflow 서버 설치:

모델 정의:

다음은 간단한 MNIST 분류 모델의 예시입니다.

import torch
from torch.nn import Linear, ReLU, CrossEntropyLoss
from torch.optim import SGD

class MNISTModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = Linear(784, 128)
        self.relu = ReLU()
        self.fc2 = Linear(128, 10)

    def forward(self, x):
        x = x.view(-1)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

PyTorch-Lightning Trainer 정의:

다음은 모델 학습을 위한 Trainer 클래스 정의입니다.

from pytorch_lightning import Trainer, LightningModule

class MNISTLightningModule(LightningModule):
    def __init__(self):
        super().__init__()
        self.model = MNISTModel()
        self.loss = CrossEntropyLoss()

    def forward(self, x, y):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = self.loss(logits, y)
        return loss

    def configure_optimizers(self):
        return SGD(self.parameters(), lr=0.01)

Trainer 클래스에 MLflowExperiment 객체를 전달하여 MLflow 추적 서버를 설정합니다.

from mlflow.tracking import MlflowExperiment

experiment = MlflowExperiment(tracking_uri="http://localhost:5000")

trainer = Trainer(
    experiment=experiment,
    max_epochs=10,
)

모델 학습 및 가중치 저장:

Trainer 객체의 fit 메서드를 사용하여 모델을 학습합니다. 학습 과정에서 모델 가중치는 MLflow 추적 서버에 자동으로 저장됩니다.
```
trainer.fit(MNISTLightningModule())
```
MLflow 추적 서버에서 모델 가중치 확인:

참고:

이 튜토리얼에서는 기본적인 예시를 제공합니다. 실제 사용 사례에 따라 코드를 수정해야 할 수도 있습니다.
MLflow 추적 서버 설정에 대한 자세한 내용은 MLflow 공식 문서를 참고하십시오.

PyTorch-Lightning을 사용하여 모델 가중치를 MLflow 추적 서버에 저장하는 예제 코드

import torch
from torch.nn import Linear, ReLU, CrossEntropyLoss
from torch.optim import SGD

from pytorch_lightning import Trainer, LightningModule
from mlflow.tracking import MlflowExperiment

# 모델 정의
class MNISTModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = Linear(784, 128)
        self.relu = ReLU()
        self.fc2 = Linear(128, 10)

    def forward(self, x):
        x = x.view(-1)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# PyTorch-Lightning Trainer 정의
class MNISTLightningModule(LightningModule):
    def __init__(self):
        super().__init__()
        self.model = MNISTModel()
        self.loss = CrossEntropyLoss()

    def forward(self, x, y):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = self.loss(logits, y)
        return loss

    def configure_optimizers(self):
        return SGD(self.parameters(), lr=0.01)

# MLflow 추적 서버 설정
experiment = MlflowExperiment(tracking_uri="http://localhost:5000")

# 모델 학습 및 가중치 저장
trainer = Trainer(
    experiment=experiment,
    max_epochs=10,
)

trainer.fit(MNISTLightningModule())

Python에서 코드를 저장합니다.
pip install 명령을 사용하여 필요한 라이브러리를 설치합니다.
코드를 실행합니다.

PyTorch-Lightning을 사용하여 모델 가중치를 MLflow 추적 서버에 저장하는 대체 방법

PyTorch-Lightning에는 MLflowLogger라는 콜백 객체가 제공됩니다. 이 콜백 객체를 사용하여 모델 학습 과정 중에 발생하는 다양한 정보를 MLflow 추적 서버에 기록할 수 있습니다. 모델 가중치도 이 정보에 포함됩니다.

from pytorch_lightning.loggers import MLflowLogger

logger = MLflowLogger(experiment_name="my-experiment")

trainer = Trainer(
    logger=logger,
    max_epochs=10,
)

trainer.fit(MNISTLightningModule())

save_checkpoint 메서드 사용:

Trainer 클래스는 save_checkpoint 메서드를 제공합니다. 이 메서드를 사용하여 모델 가중치를 직접 MLflow 추적 서버에 저장할 수 있습니다.

def save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
    # 모델 가중치를 MLflow 추적 서버에 저장
    mlflow.log_artifact(checkpoint["model_state_dict"], "model.ckpt")

trainer = Trainer(
    max_epochs=10,
    check_val_every_n_epoch=1,
    callbacks=[SaveCheckpointCallback()],
)

trainer.fit(MNISTLightningModule())

커스텀 콜백 객체 사용:

위의 방법들 외에도 커스텀 콜백 객체를 사용하여 모델 가중치를 MLflow 추적 서버에 저장할 수 있습니다. 콜백 객체는 on_train_end 메서드를 구현해야 하며, 이 메서드에서 모델 가중치를 MLflow 추적 서버에 저장하는 코드를 작성해야 합니다.