collie-mlops 0.1.0b0__py3-none-any.whl

collie/helper/pytorch/callback/callback.py

@@ -0,0 +1,155 @@
+from typing import Optional, List, Literal
+
+import torch
+
+from collie import trainer
+
+
+class Callback:
+
+    def on_train_start(
+        self,
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        ...
+
+    def on_train_end(
+        self,
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        ...
+
+    def on_epoch_start(
+        self,
+        epoch_step: int,
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        ...
+
+    def on_epoch_end(
+        self, 
+        epoch_step: int,
+        epoch_train_loss: float,
+        epoch_val_loss: Optional[float],
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        ...
+    
+    def on_batch_start(
+        self, 
+        batch_step: int, 
+        batch_data: torch.Tensor,
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        ...
+
+    def on_batch_end(
+        self, 
+        batch_step: int, 
+        batch_data: torch.Tensor,
+        batch_train_loss: float,
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        ...
+
+
+class _CallbackManager:
+    def __init__(self, callbacks: List[Callback]):
+        self.callbacks = callbacks
+
+    def on_train_start(
+        self,
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        
+        self._execute_callbacks(
+            "on_train_start",
+            trainer=trainer
+        )
+        
+    def on_train_end(
+        self,
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        
+        self._execute_callbacks(
+            "on_train_end",
+            trainer=trainer
+        )
+
+    def on_epoch_start(
+        self, 
+        trainer: "trainer.PytorchTrainer", 
+        epoch_step: int
+    ) -> None:
+
+        self._execute_callbacks(
+            "on_epoch_start",
+            trainer=trainer,
+            epoch_step=epoch_step
+        )
+
+    def on_epoch_end(
+        self, 
+        trainer: "trainer.PytorchTrainer", 
+        epoch_step: int,
+        epoch_train_loss: float,
+        epoch_val_loss: Optional[float],
+    ) -> None:
+         
+        self._execute_callbacks(
+            "on_epoch_end",
+            trainer=trainer,
+            epoch_step=epoch_step,
+            epoch_train_loss=epoch_train_loss,
+            epoch_val_loss=epoch_val_loss
+        )
+         
+    def on_batch_start(
+        self,
+        batch_step: int, 
+        batch_data: torch.Tensor,
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        
+        self._execute_callbacks(
+            "on_batch_start",
+            trainer=trainer,
+            batch_step=batch_step,
+            batch_data= batch_data
+        )
+
+    def on_batch_end(
+        self,
+        batch_step: int, 
+        batch_data: torch.Tensor,
+        batch_train_loss: float,
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        
+        self._execute_callbacks(
+            "on_batch_end",
+            trainer=trainer,
+            batch_step=batch_step,
+            batch_data= batch_data,
+            batch_train_loss=batch_train_loss
+        )
+
+    def _execute_callbacks(
+        self, 
+        method_name: Literal[
+            "on_train_start",
+            "on_train_end",
+            "on_epoch_start", 
+            "on_epoch_end", 
+            "on_batch_start", 
+            "on_batch_end"], 
+        *args, **kwargs
+    ) -> None:
+       
+        if self.callbacks is None:
+            return
+        
+        for callback in self.callbacks:
+            # Call the method specified by method_name on the callback
+            getattr(callback, method_name)(*args, **kwargs)

collie/helper/pytorch/callback/earlystop.py

@@ -0,0 +1,54 @@
+from typing import Optional
+
+from collie.helper.pytorch.callback.callback import Callback
+from collie import trainer
+from collie._common.utils import get_logger
+
+logger = get_logger()
+
+
+class EarlyStopping(Callback):
+
+    def __init__(
+        self, 
+        patience_on_epoch: int, 
+        delta: float = 0.0, 
+        monitor: str = "val_loss"
+    ):
+        super().__init__()
+
+        self.patience_on_epoch = patience_on_epoch
+        self.delta = delta
+        self.monitor = monitor
+        self.best_score = float('inf')  
+        self.wait = 0  
+        self.stopped_epoch = 0  
+        self.early_stop = False 
+
+    def on_epoch_end(
+        self, 
+        epoch_step: int, 
+        epoch_train_loss: float, 
+        epoch_val_loss: Optional[float], 
+        trainer: "trainer.PytorchTrainer"
+    ) -> None:
+        
+        if self.monitor == "val_loss" and epoch_val_loss is not None:
+            score = epoch_val_loss
+        elif self.monitor == "train_loss":
+            score = epoch_train_loss
+        else:
+            return  
+
+       
+        if self.best_score - score > self.delta:  
+            self.best_score = score
+            self.wait = 0  
+        else:
+            self.wait += 1  
+
+        if self.wait >= self.patience_on_epoch:
+            self.early_stop = True
+            self.stopped_epoch = epoch_step
+            trainer.should_stop = True
+            logger.info(f"Epoch {epoch_step}: Early stopping triggered (patience {self.patience_on_epoch} epochs).")

collie/helper/pytorch/callback/model_checkpoint.py

@@ -0,0 +1,100 @@
+from typing import Optional
+from glob import glob
+import os
+import shutil
+
+import torch
+
+from collie.helper.pytorch.callback.callback import Callback
+from collie._common.utils import get_logger
+from collie import trainer
+
+logger = get_logger()
+
+class ModelCheckpoint(Callback):
+
+    def __init__(self, topk_checkpoints: int):
+        super().__init__()
+        self.topk_checkpoints = topk_checkpoints
+        self._best_checkpoints = [] # save as (loss, epoch_idx, checkpoint_path)
+        
+        self.parent_dir = "./.checkpoint/"
+        if os.path.exists(self.parent_dir):
+            shutil.rmtree(self.parent_dir)
+            logger.info(f"Directory {self.parent_dir} has been removed.")
+
+    def on_epoch_end(
+        self,
+        trainer: "trainer.PytorchTrainer", 
+        epoch_step: int, 
+        epoch_train_loss: float, 
+        epoch_val_loss: Optional[float]
+    ) -> None:
+        
+        current_loss = epoch_val_loss if epoch_val_loss else epoch_train_loss
+        if self._should_save_checkpoint(current_loss):
+            self._save_checkpoint(
+                trainer=trainer,
+                epoch_step=epoch_step,
+                epoch_loss=epoch_train_loss,
+                loss_for_ckpt=current_loss
+            )
+    def on_train_end(self, trainer):
+        for ckpt in glob(f"{self.parent_dir}/*.pt"):
+            trainer.log_artifact(ckpt, "checkpoints")
+
+    def _should_save_checkpoint(self, current_loss: float) -> bool:
+        """
+        Returns a boolean indicating if the current checkpoint should be saved.
+
+        The current checkpoint should be saved if the topk_checkpoints list is not full.
+        Otherwise, the current checkpoint should be saved if its loss is better than the
+        loss of the worst checkpoint in the topk_checkpoints list.
+
+        Args:
+            current_loss (float): The loss of the current checkpoint.
+
+        Returns:
+            bool: True if the current checkpoint should be saved, False otherwise.
+        """
+        if len(self._best_checkpoints) < self.topk_checkpoints:
+            # If the topk_checkpoints list is not full, save the current checkpoint.
+            return True
+
+        worst_loss = max(self._best_checkpoints, key=lambda x: x[0])[0]
+        # If the current checkpoint has a better loss than the worst checkpoint in the
+        # topk_checkpoints list, save the current checkpoint.
+        return current_loss < worst_loss
+    
+    def _save_checkpoint(
+        self,
+        trainer: "trainer.PytorchTrainer", 
+        epoch_step: int, 
+        epoch_loss: float,
+        loss_for_ckpt: float
+    ) -> None:
+
+        checkpoint_path = self.parent_dir + f'model_epoch{epoch_step}.pt'
+
+        if not os.path.exists(self.parent_dir):
+            os.makedirs(self.parent_dir)
+
+        # https://zhuanlan.zhihu.com/p/136902153
+        checkpoint = {
+            "model_state_dict": trainer.model.state_dict(),
+            "optimizer_state_dict": trainer.optimizer.state_dict(),
+            'scheduler_state_dict': trainer.lr_scheduler.state_dict() if trainer.lr_scheduler else None,
+            "epoch": epoch_step,
+            "loss": epoch_loss
+        }
+
+        if len(self._best_checkpoints) < self.topk_checkpoints:
+            self._best_checkpoints.append((loss_for_ckpt, epoch_step, checkpoint_path))
+        else:
+            worst_idx, (worst_loss, _, worst_path) = max(enumerate(self._best_checkpoints), key=lambda x: x[1][0])
+
+            os.remove(worst_path)  
+            self._best_checkpoints[worst_idx] = (loss_for_ckpt, epoch_step, checkpoint_path)
+
+        self._best_checkpoints.sort(key=lambda x: x[0])
+        torch.save(checkpoint, f'{checkpoint_path}')

collie/helper/pytorch/model/loader.py

@@ -0,0 +1,55 @@
+from typing import Optional, Dict, Union
+from collections import OrderedDict
+
+import torch
+
+
+def load_pytorch_model_checkpoint(
+    ckpt_path: str, 
+    model: torch.nn.Module,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    lr_scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
+    is_ddp_model: bool = False
+) -> Dict[str, Union[torch.nn.Module, torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler, int]]:
+    """
+    Loads a PyTorch model checkpoint and optionally restores the optimizer and learning rate scheduler states.
+
+    Args:
+        ckpt_path (str): Path to the checkpoint file.
+        model (torch.nn.Module): The model to load the state dict into.
+        optimizer (Optional[torch.optim.Optimizer], optional): Optimizer to restore its state. Default is None.
+        lr_scheduler (Optional[torch.optim.lr_scheduler._LRScheduler], optional): Learning rate scheduler to restore its state. Default is None.
+        is_ddp_model (bool, optional): Whether the model is wrapped in DistributedDataParallel. Default is False.
+
+    Returns:
+        Dict[str, Union[torch.nn.Module, torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler, int]]: 
+            A dictionary containing the restored model, optimizer, learning rate scheduler (if any), and the epoch.
+    """
+
+    res_dict = dict()
+
+    checkpoint = torch.load(ckpt_path)
+    state_dict = checkpoint['model_state_dict']
+    
+    # Handle DDP (DistributedDataParallel) model
+    if is_ddp_model:
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            # remove 'module.' of DataParallel/DistributedDataParallel
+            name = k[7:] 
+            new_state_dict[name] = v
+        state_dict = new_state_dict
+
+    model.load_state_dict(state_dict, strict=False)
+    res_dict["model"] = model.eval()
+    res_dict["epoch"] = checkpoint.get("epoch", "")
+
+    if optimizer :
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        res_dict["optimizer"] = optimizer
+
+    if 'lr_scheduler_state_dict' in checkpoint and lr_scheduler:
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler_state_dict']) 
+        res_dict["lr_scheduler"] = lr_scheduler
+
+    return res_dict

collie/helper/pytorch/trainer.py

@@ -0,0 +1,304 @@
+from typing import (
+    Optional, 
+    Tuple, 
+    Dict,
+    Any,
+    List,
+)
+from abc import abstractmethod
+
+
+from tqdm import tqdm
+from torch.utils.data import DataLoader
+from torch import nn
+from torch.amp import GradScaler, autocast
+import torch
+
+from collie.trainer.trainer import Trainer
+from collie._common.decorator import type_checker
+from collie._common.utils import get_logger
+from collie.helper.pytorch.callback.callback import Callback, _CallbackManager
+from collie.helper.pytorch.callback.model_checkpoint import ModelCheckpoint
+from collie.helper.pytorch.callback.earlystop import EarlyStopping
+from collie.contracts.event import Event
+from collie.core.types import TrainerPayload, TransformerPayload
+
+
+#TODO: Features to develop:
+# 1. Train with multiple GPUs 
+logger = get_logger()
+
+
+class _AbstractPytorchTrainer:
+    
+    @abstractmethod
+    def configure_optimizers(
+        self
+    ) -> Tuple[
+        torch.optim.Optimizer, 
+        Optional[torch.optim.lr_scheduler._LRScheduler]
+        ]:
+        """
+        Configure the optimizers and learning rate schedulers.
+
+        Returns:
+            A tuple of two objects. The first object is the optimizers and the second is the learning rate schedulers if necessary.
+
+        Example:
+            code block :: python
+
+            def configure_optimizers(self) -> Tuple[List[torch.optim.Optimizer], List[torch.optim.lr_scheduler._LRScheduler]]:
+                optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.initial_lr)
+                scheduler = get_cosine_schedule_with_warmup(
+                    optimizer=optimizer,
+                    num_warmup_steps=self.num_warmup_steps,
+                    num_training_steps=self.num_training_steps
+                )
+
+                return optimizer, scheduler
+        """
+        raise NotImplementedError("Please implement the *configure_optimizers* method.")
+    
+    @type_checker(
+        (DataLoader,),
+        "The train_data in TransformerPayload should be type of DataLoader in PytorchTrainer."
+    )
+    def get_train_dataloader(self, event: Event[TransformerPayload]) -> DataLoader:
+        transformer_payload = event.payload
+        return transformer_payload.train_data
+
+    @type_checker(
+        (DataLoader, None),
+        "The validation_data in TransformerPayload should be type of DataLoader in PytorchTrainer."
+    )
+    def get_val_dataloader(self, event: Event[TransformerPayload]) -> DataLoader:
+        transformer_payload = event.payload
+        return transformer_payload.validation_data
+
+
+class PytorchTrainer(_AbstractPytorchTrainer):
+    def __init__(
+        self, 
+        model: nn.Module,
+        epochs: int,
+        device: Optional[str],
+        use_amp: Optional[bool],
+        topk_checkpoints: Optional[int],
+        earlystop_patience_on_epoch: Optional[int] = None,
+        accumulate_grad_batches: int = 1,
+        callbacks: Optional[List[Callback]] = None,
+    ) -> None:
+
+        super().__init__()
+
+        self.model = model
+
+        self.epochs = epochs
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.use_amp = use_amp
+        self.topk_checkpoints = 3 if topk_checkpoints is None else topk_checkpoints
+        self.earlystop_patience_on_epoch = earlystop_patience_on_epoch
+
+        DEFAULT_CALLBACK = [ModelCheckpoint(topk_checkpoints=self.topk_checkpoints)]
+        if self.earlystop_patience_on_epoch:
+            DEFAULT_CALLBACK.append(EarlyStopping(self.earlystop_patience_on_epoch, delta=0.0))
+        
+        self.callbacks = [] if callbacks is None else callbacks
+        self.callbacks = self.callbacks + DEFAULT_CALLBACK
+        self.accumulate_grad_batches = accumulate_grad_batches
+
+        self.train_data_loader = None
+        self.val_data_loader = None
+        self.optimizer = None
+        self.lr_scheduler = None
+        self.grad_scaler = None
+        self.should_stop = False
+        self.cb_manager = _CallbackManager(callbacks=self.callbacks)
+
+        self.model = self.model.to(self.device)
+
+    def run(
+        self, 
+        event: Event[TransformerPayload]
+    ) -> Event[TrainerPayload]:
+
+        self.train_data_loader = self.get_train_dataloader(event=event)
+        self.val_data_loader = self.get_val_dataloader(event=event)
+
+        self.optimizer, self.lr_scheduler = self.configure_optimizers()
+        
+        trainstep_per_epoch = len(self.train_data_loader)
+        total_train_step = trainstep_per_epoch * self.epochs
+
+        self.log_param("epoch", self.epochs)
+        self.log_param("batch size", self.train_data_loader.batch_size)
+        self.log_param("total training step", total_train_step)
+
+        self.cb_manager.on_train_start(self)
+        self.optimizer.zero_grad()
+        # TODO: the train loop is written by the user in handle method.
+        for epoch_idx in tqdm(range(1, self.epochs + 1)):
+            # Modified by Earlystop callback
+            if self.should_stop:
+                break
+
+            self.cb_manager.on_epoch_start(self, epoch_idx)
+            running_loss = 0
+            self.model.train()
+            for batch_idx, batch_data in enumerate(self.train_data_loader, start=1):
+                self.cb_manager.on_batch_start(
+                    batch_data=batch_data,
+                    batch_step=batch_idx,
+                    trainer=self
+                )
+                
+                loss = self._get_train_step_result(epoch_idx, batch_data) / self.accumulate_grad_batches
+                running_loss += loss.item()  * self.accumulate_grad_batches
+
+                self._backward(loss=loss, batch_idx=batch_idx)
+                # The learning rate scheduler is batch level.
+                if self.lr_scheduler and batch_idx % self.accumulate_grad_batches == 0:
+                    self.lr_scheduler.step()
+
+                self.cb_manager.on_batch_end(
+                    batch_step=batch_idx,
+                    batch_data=batch_data,
+                    batch_train_loss=running_loss,
+                    trainer=self
+                )
+
+            epoch_loss = running_loss / trainstep_per_epoch
+            
+            self.log_metric("train loss", epoch_loss, step=epoch_idx)
+            self._log_lr(epoch_step=epoch_idx)
+
+            epoch_val_loss = self.validation_loop(epoch_step=epoch_idx)
+            
+            self.cb_manager.on_epoch_end(
+                trainer=self,
+                epoch_step=epoch_idx,
+                epoch_train_loss=epoch_loss,
+                epoch_val_loss=epoch_val_loss
+            )
+        
+        # TODO: Fix the issue of loading the model from the following result model.
+        self.log_model(self.model, model_type="pt")
+
+        self.cb_manager.on_train_end(self)
+
+        # return {"model":self.model, "loss": epoch_loss}
+        trainer_payload = TrainerPayload(
+            model=self.model,
+            train_loss=epoch_loss,
+            val_loss=epoch_val_loss
+        )
+        return Event(
+            payload=trainer_payload
+        )
+
+    def validation_loop(self, epoch_step: int) -> Optional[float]:
+        
+        if self.val_data_loader is None:
+            return None
+        
+        valstep_per_epoch = len(self.val_data_loader)
+
+        self.model.eval()
+        val_running_loss = 0
+        with torch.no_grad():
+            
+            for val_batch_data in self.val_data_loader:
+                
+                val_loss = self.validation_step(epoch_step, val_batch_data)
+                val_running_loss += val_loss.item()
+
+            epoch_val_loss = val_running_loss / valstep_per_epoch
+        
+        self.log_metric("val_loss", epoch_val_loss, step=epoch_step)
+        return epoch_val_loss
+
+    def _backward(
+        self, 
+        loss: torch.Tensor, 
+        batch_idx: int
+    ) -> None:
+
+        def should_step_optimizer(batch_idx: int) -> bool:
+            return batch_idx % self.accumulate_grad_batches == 0 or batch_idx == len(self.train_data_loader)
+
+        if self.use_amp:
+            if self.grad_scaler is None:
+                self.grad_scaler = GradScaler(device=self.device)
+
+            self.grad_scaler.scale(loss).backward()
+            if should_step_optimizer(batch_idx):
+                
+                if self._has_invalid_gradients():
+                    # gradient clipping
+                    self.grad_scaler.unscale_(self.optimizer)
+                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+                
+                self.grad_scaler.step(self.optimizer)  
+                self.grad_scaler.update() 
+                self.optimizer.zero_grad()
+        else:
+            loss.backward()
+            if should_step_optimizer():
+                if self._has_invalid_gradients(batch_idx):
+                    # gradient clipping
+                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+                self.optimizer.step()
+                self.optimizer.zero_grad()
+    
+    def _has_invalid_gradients(self) -> bool:
+   
+        for param in self.model.parameters():
+            if param.grad is not None:
+                if torch.isnan(param.grad).any() or torch.isinf(param.grad).any():
+                    logger.warning(f"Detected NaN or Inf gradients for parameter {param}. Zeroing out the gradients.")
+                    # side effect
+                    param.grad.zero_()
+                    return True
+        return False
+            
+    def _log_lr(self, epoch_step: int) -> None:
+        
+        lr = self.optimizer.param_groups[0]['lr'] if not self.lr_scheduler else self.lr_scheduler.get_last_lr()[0]
+        self.log_metric("learning rate", lr, step=epoch_step)
+ 
+    @type_checker((torch.Tensor,), 
+                  "The return type of *train_step* method must be 'Tensor'."
+    )
+    def _get_train_step_result(
+        self, 
+        epoch_step: int, 
+        batch_data: torch.Tensor
+    ):
+
+        if self.use_amp:
+            with autocast(device_type=self.device):
+                loss = self.handle(epoch_step, batch_data)
+        else:
+            loss = self.handle(epoch_step, batch_data)
+
+        return loss
+    
+    def _get_optimizers(
+        self
+    ) -> Tuple[torch.optim.Optimizer, Optional[torch.optim.lr_scheduler.LRScheduler]]:
+        """
+        Get the optimizer and learning rate scheduler from the 
+        *configure_optimizers* method.
+        """
+        result = self.configure_optimizers()
+        if isinstance(result, tuple):
+            optimizers, lr_scheduler = result
+            if not isinstance(lr_scheduler, torch.optim.lr_scheduler._LRScheduler):
+                raise TypeError(
+                    f"learning rate scheduler must be type of \
+                    torch.optim.lr_scheduler._LRScheduler "
+                )
+        else:
+            optimizers, lr_scheduler = result, None
+        return optimizers, lr_scheduler
+