dinhdat1110/diffusion-model

1import os2import diffusion3from pytorch_lightning.callbacks import (4    ModelCheckpoint,5    LearningRateMonitor6)7 8 9class ModelCallback:10    def __init__(11        self,12        root_path: str,13        ckpt_monitor: str = "val_loss",14        ckpt_mode: str = "min",15    ):16        ckpt_path = os.path.join(os.path.join(root_path, "model/"))17        if not os.path.exists(root_path):18            os.makedirs(root_path)19        if not os.path.exists(ckpt_path):20            os.makedirs(ckpt_path)21 22        self.ckpt_callback = ModelCheckpoint(23            monitor=ckpt_monitor,24            dirpath=ckpt_path,25            filename="model",26            save_top_k=1,27            mode=ckpt_mode,28            save_weights_only=True29        )30 31        self.lr_callback = LearningRateMonitor("step")32 33        self.ema_callback = diffusion.EMACallback(decay=0.995)34 35    def get_callback(self):36        return [37            self.ckpt_callback, self.lr_callback, self.ema_callback38        ]39