Jae-Kyung Cho LLM Developers who was a Robotics engineer

MLOps study - Raviraja Week 1: W&B

Week 1은 Weight & Bias 사용법에 대한 내용입니다. ICRA 워크숍에 다녀왔을 때, 그리고 Google brain 친구들과 협업할 때 사용하는 방법을 배웠던 툴들인데, 확실히 아주 쓰기 편합니다. 딥러닝에 대한 이론을 공부하고 논문을 작성하는 것도 중요하지만 발빠르게 좋은 툴을 익히고 내것으로 만드는 것도 중요하다는 생각이 들게 되었습니다.
WandB에서 사용할 수 있는 기능은 크게 3 가지가 있습니다. 3 가지 기능을 사용하기 전에 configuration 부터 시작해 보겠습니다.




WandB configuring

WandB는 Pytorch lightning 과 호환성이 굉장히 좋아 보입니다. 물론 Pytorch와의 연동성도 좋았지만 API를 가져다가 직접 logging해야 했다면, Pytorch lightning은 처음 Logger 선언만으로 많은 기능을 사용할 수 있었습니다. 선언시 사용할 수 있는 Argument들은 다음과 같습니다.

< Arguments >

  • name: Display name for the run.
  • save_dir: Path where data is saved.
  • offline: Run offline (data can be streamed later to wandb servers).
  • id: Sets the version, mainly used to resume a previous run.
  • version: Same as id.
  • anonymous: Enables or explicitly disables anonymous logging.
  • project: The name of the project to which this run will belong.
  • log_model: Log checkpoints created by ~pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint as W&B artifacts. latest and best aliases are automatically set.
    • log_model=’all’: save all checkpoints
    • log_model=True: save last checkpoint
from pytorch_lightning.loggers import WandbLogger
wandb_logger = WandbLogger(project="MLOps Basics")

trainer = pl.Trainer(
        max_epochs=3,
        logger=wandb_logger,
        callbacks=[],
)




WandB logging

torchmetrics를 사용하면 여러가지 metric들을 편하게 계산해 wandb logger에 저장할 수 있습니다. 간단하게 self.log 접근 가능하다는 것이 Pytorch Lightning의 장점인 듯 합니다!!

logging을 할 때에는 여러가지 옵션을 사용할 수 있습니다.

  • prog_bar=True : progress bar를 보여줌
  • on_epoch=True : Epoch 내의 batches에 대한 평균 값을 취함
  • on_step=True : 모든 batch에 대한 값을 logging -> validation 시에는 False로 하는 것이 일반적이다.

LightningModule의 method 별로 logging 하는 주기가 다릅니다.

  • training_step: 각 batch 마다 logging
  • training_epoch_end: Epoch 마다 logging
  • validation_step: 각 batch 마다 logging
  • validation_epoch_end: Epoch 마다 logging

epoch_end 에서는 각 step의 결과를 가지고 Epoch 단위로 합치는 작업을 수행하기 때문에, step에서는 dictionary의 형태로 결과를 전달해야 합니다.

class ColaModel(pl.LightningModule):
    def __init__(self, model_name="google/bert_uncased_L-2_H-128_A-2", lr=3e-5):
        self.train_accuracy_metric = torchmetrics.Accuracy()
        self.val_accuracy_metric = torchmetrics.Accuracy()

    def training_step(self, batch, batch_idx):
        outputs = self.forward(
            batch["input_ids"], batch["attention_mask"], labels=batch["label"]
        )
        preds = torch.argmax(outputs.logits, 1)
        train_acc = self.train_accuracy_metric(preds, batch["label"])

        # self.log로 wandb logger에 접근 가능하다. 따로 선언 안해줘도 됨.
        self.log("train/loss", outputs.loss, prog_bar=True, on_epoch=True)
        self.log("train/acc", train_acc, prog_bar=True, on_epoch=True)
        return outputs.loss
    
    def validation_step(self, batch, batch_idx):
        labels = batch["label"]
        outputs = self.forward(
            batch["input_ids"], batch["attention_mask"], labels=batch["label"]
        )
        preds = torch.argmax(outputs.logits, 1)
        valid_acc = self.val_accuracy_metric(preds, labels)

        self.log("valid/loss", outputs.loss, prog_bar=True, on_step=True)
        return {"labels": labels, "logits": outputs.logits}

    def validation_epoch_end(self, outputs):
        # validation_step의 결과 값을 가지고 Epoch 마다 새로 계산하는 것
        labels = torch.cat([x["labels"] for x in outputs])
        logits = torch.cat([x["logits"] for x in outputs])
        preds = torch.argmax(logits, 1)

        data = confusion_matrix(labels.numpy(), preds.numpy())
        return data




WandB plotting

Plot은 여러 가지로 가능한데, 가장 마음에 드는 방식은 matplotlib와 연동하는 방식입니다. 기존에 사용하던대로 plt로 plot한 것을 wandb에 넣는 것만으로도 visualization이 가능합니다.

def validation_epoch_end(self, outputs):
    # validation_step의 결과 값을 가지고 Epoch 마다 새로 계산하는 것
    labels = torch.cat([x["labels"] for x in outputs])
    logits = torch.cat([x["logits"] for x in outputs])
    preds = torch.argmax(logits, 1)

    data = confusion_matrix(labels.numpy(), preds.numpy())
    df_cm = pd.DataFrame(data, columns=np.unique(labels), index=np.unique(labels))
    df_cm.index.name = "Actual"
    df_cm.columns.name = "Predicted"
    plt.figure(figsize=(7, 4))
    plot = sns.heatmap(
        df_cm, cmap="Blues", annot=True, annot_kws={"size": 16}
    )  # font size
    self.logger.experiment.log({"Confusion Matrix": wandb.Image(plot)})

plt로 그린 그림을 self.logger.experiment.log 내부에 wandb.Image로 변환해서 넣어주기만 하면 된다!! 엄청난 기술의 발전이구만~~!~!~!




Keep watching data samples

모델을 학습시키는 과정에서 가장 궁금한 것은, 어떤 sample에서 성공적으로 동작하고, 어떤 sample에서 제대로 동작하지 않는지 입니다.. 어디서 안되고 어디서 되는지 알아야, overfitting의 여부를 판별하고 모델의 개선점을 파악할 수 있기 때문이죠. 이 또한 callback logger 가 존재합니다!! (엄청나군 정말…)

on_validation_end 메쏘드에서 Table 형식으로 logging 한다면 더 쉽게 파악이 가능합니다. 물론 코딩을 조금 하긴 해야 함!!

class SamplesVisualisationLogger(pl.Callback):
    def __init__(self, datamodule):
        super().__init__()
        self.datamodule = datamodule

    def on_validation_end(self, trainer, pl_module):
        # can be done on complete dataset also
        val_batch = next(iter(self.datamodule.val_dataloader()))
        sentences = val_batch["sentence"]

        # get the predictions
        outputs = pl_module(val_batch["input_ids"], val_batch["attention_mask"])
        preds = torch.argmax(outputs.logits, 1)
        labels = val_batch["label"]

        # predicted and labelled data
        df = pd.DataFrame(
            {"Sentence": sentences, "Label": labels.numpy(), "Predicted": preds.numpy()}
        )

        # wrongly predicted data
        wrong_df = df[df["Label"] != df["Predicted"]]

        # Logging wrongly predicted dataframe as a table
        trainer.logger.experiment.log(
            {
                "examples": wandb.Table(dataframe=wrong_df, allow_mixed_types=True),
                "global_step": trainer.global_step,
            }
        )

ipynb 파일 다운로드




references:

Comments