MLOps study - Raviraja Week 0: Pytorch Lightning
11 Oct 2022Raviraja의 블로그 포스트들을 참고하여 MLOps에 대해 공부해 보려 합니다. 시작이 반이다!! :D
Week 0는 MLOps에 대해 공부하기 위한 환경을 세팅합니다. Raviraja는 NLP 쪽을 연구하기 때문에 환경 설정이 NLP, 그리고 약간의 classification에 치중되어 있다는 점을 인지하고 넘어가겠습니다. (추후 RL을 이용하는 MLOps로 발전시켜 나가겠습니다.) 기본적으로 Pytorch-Lightning 라이브러리를 사용합니다. Pytorch lightning은 pytorch wrapper의 일종입니다 :D
Pytorch Lightning은 크게 4가지 모듈로 이루어집니다. 차례로 살펴 보도록 하겠습니다.
DataModule
Pytorch lightning의 경우 Pytorch의 DataLoader와 유사한 DataModule을 사용합니다. DataLoader를 사용하기 전에 데이터를 전처리하는 과정이 있는데, 그걸 다 모듈 안에 포함시킨 거라고 보면 되겠습니다.
< 정의해야 하는 method >
- prepare_data -> 데이터 다운로드
- setup -> 데이터 전처리
- train_dataloader, val_dataloader, test_dataloader -> 데이터 로더
< DataModule 안에서 수행되는 작업 >
- Download / tokenize / process
- Clean and save to disk
- Load inside Dataset
- Apply transforms (rotate, tokenize, etc…)
- Wrap inside a DataLoader (Pytorch)
class DataModule(pl.LightningDataModule):
def __init__(self, model_name="google/bert_uncased_L-2_H-128_A-2", batch_size=32):
super().__init__()
self.batch_size = batch_size
self.tokenizer = AutoTokenizer.from_pretrained(model_name) # Transformer (BERT) model
def prepare_data(self):
cola_dataset = load_dataset("glue", "cola")
self.train_data = cola_dataset["train"]
self.val_data = cola_dataset["validation"]
def tokenize_data(self, example):
# processing the data
return self.tokenizer(
example["sentence"],
truncation=True,
padding="max_length",
max_length=256,
)
def setup(self, stage=None):
if stage == "fit" or stage is None:
self.train_data = self.train_data.map(self.tokenize_data, batched=True)
self.train_data.set_format(
type="torch", columns=["input_ids", "attention_mask", "label"]
)
self.val_data = self.val_data.map(self.tokenize_data, batched=True)
self.val_data.set_format(
type="torch", columns=["input_ids", "attention_mask", "label"]
)
def train_dataloader(self):
return torch.utils.data.DataLoader(
self.train_data, batch_size=self.batch_size, shuffle=True
)
def val_dataloader(self):
return torch.utils.data.DataLoader(
self.val_data, batch_size=self.batch_size, shuffle=False
)
LightningModule
Pytorch에서 model을 만들떄 상속받았던 torch.nn.Module과 마찬가지로 Pytorch-lightning은 pl.LightningModule을 상속받습니다. forward 만 정의해 주면 되었던 때와는 다르게, 몇 가지 method를 추가로 정의해 주어야 합니다. (Document)
< 정의 해야 할 methods >
- forward -> 모델 forward
- training_step -> Update 및 Loss 계산
- validation_step
- test_step (optional)
- configure_optimizers -> Optimizer 초기화
class ColaModel(pl.LightningModule):
def __init__(self, model_name="google/bert_uncased_L-2_H-128_A-2", lr=1e-2):
super(ColaModel, self).__init__()
self.save_hyperparameters()
self.bert = AutoModel.from_pretrained(model_name)
self.W = nn.Linear(self.bert.config.hidden_size, 2)
self.num_classes = 2
def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
h_cls = outputs.last_hidden_state[:, 0]
logits = self.W(h_cls)
return logits
def training_step(self, batch, batch_idx):
logits = self.forward(batch["input_ids"], batch["attention_mask"])
loss = F.cross_entropy(logits, batch["label"])
self.log("train_loss", loss, prog_bar=True)
return loss
def validation_step(self, batch, batch_idx):
logits = self.forward(batch["input_ids"], batch["attention_mask"])
loss = F.cross_entropy(logits, batch["label"])
_, preds = torch.max(logits, dim=1)
val_acc = accuracy_score(preds.cpu(), batch["label"].cpu())
val_acc = torch.tensor(val_acc)
self.log("val_loss", loss, prog_bar=True)
self.log("val_acc", val_acc, prog_bar=True)
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=self.hparams["lr"])
Trainer
DataModule과 Pytorch-lightning model은 Trainer를 이용해서 학습을 진행하게 됩니다. Tensorflow의 Session과 비슷한 접근 방법이라고 볼 수 있겠습니다.
< Trainer 가 사용할 수 있는 options 예시 >
- logging
- gradient accumulation
- half precision training
- distributed computing
< Loggers >
- TensorboardLogger
- WandbLogger
< Callbacks > Documents
cola_data = DataModule()
cola_model = ColaModel()
checkpoint_callbacks = [
ModelCheckpoint(dirpath="./models", monitor="val_loss", mode="min"), # Save model
EarlyStopping(monitor="val_loss", patience=3, verbose=True, mode="min"),
]
trainer = pl.Trainer(
gpus=(1 if torch.cuda.is_available() else 0),
max_epochs=1,
fast_dev_run=False, # True: one batch training one validation -> for debugging
logger=pl.loggers.TensorBoardLogger("logs/", name="cola", version=1), # directory: logs/cola
# logger = pl.loggers.WandbLogger(name='cola',project='pytorchlightning')
callbacks=checkpoint_callbacks,
)
trainer.fit(cola_model, cola_data)
Inference
MLOps는 모델의 Training과 Inference의 모듈을 분리합니다. 서버에서 학습이 진행되는 동안에도 모델을 freeze 하고 버전 관리를 하며 debuggin 할 수 있어야 하기 때문입니다.
< 정의 해야 할 methods >
- predict
< Inference 내부에서 수행되는 작업 >
- Load the trained model
- Get the input
- Convert the input in the required format
- Get the predictions
class ColaPredictor:
def __init__(self, model_path):
self.model_path = model_path
# loading the trained model
self.model = ColaModel.load_from_checkpoint(model_path)
# keep the model in eval mode
self.model.eval()
self.model.freeze()
self.processor = DataModule()
self.softmax = torch.nn.Softmax(dim=0)
self.lables = ["unacceptable", "acceptable"]
def predict(self, text):
# text => run time input
inference_sample = {"sentence": text}
# tokenizing the input
processed = self.processor.tokenize_data(inference_sample)
# predictions
logits = self.model(
torch.tensor([processed["input_ids"]]),
torch.tensor([processed["attention_mask"]]),
)
scores = self.softmax(logits[0]).tolist()
predictions = []
for score, label in zip(scores, self.lables):
predictions.append({"label": label, "score": score})
return predictions
사실 그렇게 대단한 변화가 있어보이지는 않습니다만, Pytorch 가 아이스크림이라면 Pytorch Lightning은 Cherry on top 이라고 합니다. 아직 MLOps에 해당되는 기능은 어떤 것이 있는지 잘 모르겠지만, Pytorch와의 호환성을 생각한다면, 훨씬 간단하게 기능들을 사용할 수 있을 것 같습니다 :ㅇ