libcom 0.1.8__py3-none-any.whl

libcom/image_harmonization/source/src/lbm/trainer/loggers.py

@@ -0,0 +1,324 @@
+import logging
+import math
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+import torch
+import wandb
+from PIL import Image, ImageDraw, ImageFont
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import Callback
+from pytorch_lightning.utilities import rank_zero_only
+from torchvision.utils import make_grid
+
+from ..trainer import TrainingPipeline
+
+logging.basicConfig(level=logging.INFO)
+
+
+def create_grid_texts(
+    texts: List[str],
+    n_cols: int = 4,
+    image_size: Tuple[int] = (512, 512),
+    font_size: int = 40,
+    margin: int = 5,
+    offset: int = 5,
+) -> Image.Image:
+    """
+    Create a grid of white images containing the given texts.
+
+    Args:
+        texts (List[str]): List of strings to be drawn on images.
+        n_cols (int): Number of columns in the grid.
+        image_size (tuple): Size of the generated images (width, height).
+        font_size (int): Font size of the text.
+        margin (int): Margin around the text.
+        offset (int): Offset between lines.
+
+    Returns:
+        PIL.Image: List of generated images as a grid
+    """
+
+    images = []
+    font = ImageFont.load_default(size=font_size)
+
+    for text in texts:
+        img = Image.new("RGB", image_size, color="white")
+        draw = ImageDraw.Draw(img)
+        margin_ = margin
+        offset_ = offset
+        for line in wrap_text(
+            text=text, draw=draw, max_width=image_size[0] - 2 * margin_, font=font
+        ):
+            draw.text((margin_, offset_), line, font=font, fill="black")
+            offset_ += font_size
+        images.append(img)
+
+    # create a pil grid
+    n_rows = math.ceil(len(images) / n_cols)
+    grid = Image.new(
+        "RGB", (n_cols * image_size[0], n_rows * image_size[1]), color="white"
+    )
+    for i, img in enumerate(images):
+        grid.paste(img, (i % n_cols * image_size[0], i // n_cols * image_size[1]))
+
+    return grid
+
+
+def wrap_text(
+    text: str, draw: ImageDraw.Draw, max_width: int, font: ImageFont
+) -> List[str]:
+    """
+    Wrap text to fit within a specified width when drawn.
+    It will return to the new line when the text is larger than the max_width.
+
+    Args:
+        text (str): The text to be wrapped.
+        draw (ImageDraw.Draw): The draw object to calculate text size.
+        max_width (int): The maximum width for the wrapped text.
+        font (ImageFont): The font used for the text.
+
+    Returns:
+        List[str]: List of wrapped lines.
+    """
+    lines = []
+    current_line = ""
+    for letter in text:
+        if draw.textbbox((0, 0), current_line + letter, font=font)[2] <= max_width:
+            current_line += letter
+        else:
+            lines.append(current_line)
+            current_line = letter
+    lines.append(current_line)
+    return lines
+
+
+class WandbSampleLogger(Callback):
+    """
+    Logger for logging samples to wandb. This logger is used to log images, text, and metrics to wandb.
+
+    Args:
+        log_batch_freq (int): The frequency of logging samples to wandb. Default is 100.
+    """
+
+    def __init__(self, log_batch_freq: int = 100):
+        super().__init__()
+        self.log_batch_freq = log_batch_freq
+
+    def on_train_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: TrainingPipeline,
+        outputs: Dict[str, Any],
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        self.log_samples(trainer, pl_module, outputs, batch, batch_idx, split="train")
+        self._process_logs(trainer, outputs, split="train")
+
+    def on_validation_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: TrainingPipeline,
+        outputs: Dict[str, Any],
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        self.log_samples(trainer, pl_module, outputs, batch, batch_idx, split="val")
+        self._process_logs(trainer, outputs, split="val")
+
+    @rank_zero_only
+    @torch.no_grad()
+    def log_samples(
+        self,
+        trainer: Trainer,
+        pl_module: TrainingPipeline,
+        outputs: Dict[str, Any],
+        batch: Dict[str, Any],
+        batch_idx: int,
+        split: str = "train",
+    ) -> None:
+        if hasattr(pl_module, "log_samples"):
+            if batch_idx % self.log_batch_freq == 0:
+                is_training = pl_module.training
+                if is_training:
+                    pl_module.eval()
+
+                logs = pl_module.log_samples(batch)
+                logs = self._process_logs(trainer, logs, split=split)
+
+                if is_training:
+                    pl_module.train()
+        else:
+            logging.warning(
+                "log_img method not found in LightningModule. Skipping image logging."
+            )
+
+    @rank_zero_only
+    def _process_logs(
+        self, trainer, logs: Dict[str, Any], rescale=True, split="train"
+    ) -> Dict[str, Any]:
+        for key, value in logs.items():
+            if isinstance(value, torch.Tensor):
+                value = value.detach().cpu()
+                if value.dim() == 4:
+                    images = value
+                    if rescale:
+                        images = (images + 1.0) / 2.0
+                    grid = make_grid(images, nrow=4)
+                    grid = grid.permute(1, 2, 0)
+                    grid = grid.mul(255).clamp(0, 255).to(torch.uint8)
+                    logs[key] = grid.numpy()
+                    trainer.logger.experiment.log(
+                        {f"{key}/{split}": [wandb.Image(Image.fromarray(logs[key]))]},
+                        step=trainer.global_step,
+                    )
+
+                # Scalar tensor
+                if value.dim() == 1 or value.dim() == 0:
+                    value = value.float().numpy()
+                    trainer.logger.experiment.log(
+                        {f"{key}/{split}": value}, step=trainer.global_step
+                    )
+
+            # list of string (e.g. text)
+            if isinstance(value, list):
+                if isinstance(value[0], str):
+                    pil_image_texts = create_grid_texts(value)
+                    wandb_image = wandb.Image(pil_image_texts)
+                    trainer.logger.experiment.log(
+                        {f"{key}/{split}": [wandb_image]},
+                        step=trainer.global_step,
+                    )
+
+            # dict of tensors (e.g. metrics)
+            if isinstance(value, dict):
+                for k, v in value.items():
+                    if isinstance(v, torch.Tensor):
+                        value[k] = v.detach().cpu().numpy()
+                trainer.logger.experiment.log(
+                    {f"{key}/{split}": value}, step=trainer.global_step
+                )
+
+            if isinstance(value, int) or isinstance(value, float):
+                trainer.logger.experiment.log(
+                    {f"{key}/{split}": value}, step=trainer.global_step
+                )
+
+        return logs
+
+
+class TensorBoardSampleLogger(Callback):
+    """
+    Logger for logging samples to tensorboard. This logger is used to log images, text, and metrics to tensorboard.
+
+    Args:
+        log_batch_freq (int): The frequency of logging samples to tensorboard. Default is 100.
+    """
+
+    def __init__(self, log_batch_freq: int = 100):
+        super().__init__()
+        self.log_batch_freq = log_batch_freq
+
+    def on_train_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: TrainingPipeline,
+        outputs: Dict[str, Any],
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        self.log_samples(trainer, pl_module, outputs, batch, batch_idx, split="train")
+        self._process_logs(trainer, outputs, split="train")
+
+    def on_validation_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: TrainingPipeline,
+        outputs: Dict[str, Any],
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        self.log_samples(trainer, pl_module, outputs, batch, batch_idx, split="val")
+        self._process_logs(trainer, outputs, split="val")
+
+    @rank_zero_only
+    @torch.no_grad()
+    def log_samples(
+        self,
+        trainer: Trainer,
+        pl_module: TrainingPipeline,
+        outputs: Dict[str, Any],
+        batch: Dict[str, Any],
+        batch_idx: int,
+        split: str = "train",
+    ) -> None:
+        if hasattr(pl_module, "log_samples"):
+            if batch_idx % self.log_batch_freq == 0:
+                is_training = pl_module.training
+                if is_training:
+                    pl_module.eval()
+
+                logs = pl_module.log_samples(batch)
+                logs = self._process_logs(trainer, logs, split=split)
+
+                if is_training:
+                    pl_module.train()
+        else:
+            logging.warning(
+                "log_img method not found in LightningModule. Skipping image logging."
+            )
+
+    @rank_zero_only
+    def _process_logs(
+        self, trainer, logs: Dict[str, Any], rescale=True, split="train"
+    ) -> Dict[str, Any]:
+        for key, value in logs.items():
+            if isinstance(value, torch.Tensor):
+                value = value.detach().cpu()
+                if value.dim() == 4:
+                    images = value
+                    if rescale:
+                        images = (images + 1.0) / 2.0
+                    grid = make_grid(images, nrow=4)
+                    # grid = grid.permute(1, 2, 0)
+                    grid = grid.mul(255).clamp(0, 255).to(torch.uint8)
+                    logs[key] = grid.numpy()
+                    trainer.logger.experiment.add_image(
+                        f"{key}/{split}",
+                        logs[key],
+                        trainer.global_step,
+                    )
+
+                # Scalar tensor
+                if value.dim() == 1 or value.dim() == 0:
+                    value = value.float().numpy()
+                    trainer.logger.experiment.add_scalar(
+                        f"{key}/{split}", value, trainer.global_step
+                    )
+
+            # list of string (e.g. text)
+            if isinstance(value, list):
+                if isinstance(value[0], str):
+                    pil_image_texts = create_grid_texts(value)
+                    trainer.logger.experiment.add_image(
+                        f"{key}/{split}",
+                        np.transpose(np.array(pil_image_texts), (2, 0, 1)),
+                        trainer.global_step,
+                    )
+
+            # dict of tensors (e.g. metrics)
+            if isinstance(value, dict):
+                for k, v in value.items():
+                    if isinstance(v, torch.Tensor):
+                        value[k] = v.detach().cpu().numpy()
+                trainer.logger.experiment.add_scalar(
+                    f"{key}/{split}", value, trainer.global_step
+                )
+
+            if isinstance(value, int) or isinstance(value, float):
+                trainer.logger.experiment.add_scalar(
+                    f"{key}/{split}", value, trainer.global_step
+                )
+
+        return logs

libcom/image_harmonization/source/src/lbm/trainer/trainer.py

@@ -0,0 +1,199 @@
+import importlib
+import logging
+import re
+import time
+from typing import Any, Dict
+
+import pytorch_lightning as pl
+import torch
+
+from ..models.base.base_model import BaseModel
+from .training_config import TrainingConfig
+
+logging.basicConfig(level=logging.INFO)
+
+
+class TrainingPipeline(pl.LightningModule):
+    """
+    Main Training Pipeline class
+
+    Args:
+
+        model (BaseModel): The model to train
+        pipeline_config (TrainingConfig): The configuration for the training pipeline
+        verbose (bool): Whether to print logs in the console. Default is False.
+    """
+
+    def __init__(
+        self,
+        model: BaseModel,
+        pipeline_config: TrainingConfig,
+        verbose: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.model = model
+        self.pipeline_config = pipeline_config
+        self.log_samples_model_kwargs = pipeline_config.log_samples_model_kwargs
+
+        # save hyperparameters.
+        self.save_hyperparameters(ignore="model")
+        self.save_hyperparameters({"model_config": model.config.to_dict()})
+
+        # logger.
+        self.verbose = verbose
+
+        # setup logging.
+        log_keys = pipeline_config.log_keys
+
+        if isinstance(log_keys, str):
+            log_keys = [log_keys]
+
+        if log_keys is None:
+            log_keys = []
+
+        self.log_keys = log_keys
+
+    def on_fit_start(self) -> None:
+        self.model.on_fit_start(device=self.device)
+        if self.global_rank == 0:
+            self.timer = time.perf_counter()
+
+    def on_train_batch_end(
+        self, outputs: Dict[str, Any], batch: Any, batch_idx: int
+    ) -> None:
+        if self.global_rank == 0:
+            logging.debug("on_train_batch_end")
+        self.model.on_train_batch_end(batch)
+
+        average_time_frequency = 10
+        if self.global_rank == 0 and batch_idx % average_time_frequency == 0:
+            delta = time.perf_counter() - self.timer
+            logging.info(
+                f"Average time per batch {batch_idx} took {delta / (batch_idx + 1)} seconds"
+            )
+
+    def configure_optimizers(self) -> torch.optim.Optimizer:
+        """
+        Setup optimizers and learning rate schedulers.
+        """
+        optimizers = []
+        lr = self.pipeline_config.learning_rate
+        param_list = []
+        n_params = 0
+        param_list_ = {"params": []}
+        for name, param in self.model.named_parameters():
+            for regex in self.pipeline_config.trainable_params:
+                pattern = re.compile(regex)
+                if re.match(pattern, name):
+                    if param.requires_grad:
+                        param_list_["params"].append(param)
+                        n_params += param.numel()
+
+        param_list.append(param_list_)
+
+        logging.info(f"Number of trainable parameters: {n_params}")
+
+        optimizer_cls = getattr(
+            importlib.import_module("torch.optim"),
+            self.pipeline_config.optimizer_name,
+        )
+        optimizer = optimizer_cls(
+            param_list, lr=lr, **self.pipeline_config.optimizer_kwargs
+        )
+        optimizers.append(optimizer)
+
+        self.optims = optimizers
+        schedulers_config = self.configure_lr_schedulers()
+
+        for name, param in self.model.named_parameters():
+            set_grad_false = True
+            for regex in self.pipeline_config.trainable_params:
+                pattern = re.compile(regex)
+                if re.match(pattern, name):
+                    if param.requires_grad:
+                        set_grad_false = False
+            if set_grad_false:
+                param.requires_grad = False
+
+        num_trainable_params = sum(
+            p.numel() for p in self.model.parameters() if p.requires_grad
+        )
+
+        logging.info(f"Number of trainable parameters: {num_trainable_params}")
+
+        schedulers_config = self.configure_lr_schedulers()
+
+        if schedulers_config is None:
+            return optimizers
+
+        return optimizers, [
+            schedulers_config_ for schedulers_config_ in schedulers_config
+        ]
+
+    def configure_lr_schedulers(self):
+        schedulers_config = []
+        if self.pipeline_config.lr_scheduler_name is None:
+            scheduler = None
+            schedulers_config.append(scheduler)
+        else:
+            scheduler_cls = getattr(
+                importlib.import_module("torch.optim.lr_scheduler"),
+                self.pipeline_config.lr_scheduler_name,
+            )
+            scheduler = scheduler_cls(
+                self.optims[0],
+                **self.pipeline_config.lr_scheduler_kwargs,
+            )
+            lr_scheduler_config = {
+                "scheduler": scheduler,
+                "interval": self.pipeline_config.lr_scheduler_interval,
+                "monitor": "val_loss",
+                "frequency": self.pipeline_config.lr_scheduler_frequency,
+            }
+            schedulers_config.append(lr_scheduler_config)
+
+        if all([scheduler is None for scheduler in schedulers_config]):
+            return None
+
+        return schedulers_config
+
+    def training_step(self, train_batch: Dict[str, Any], batch_idx: int) -> dict:
+        model_output = self.model(train_batch)
+        loss = model_output["loss"]
+        logging.info(f"loss: {loss}")
+        return {
+            "loss": loss,
+            "batch_idx": batch_idx,
+        }
+
+    def validation_step(self, val_batch: Dict[str, Any], val_idx: int) -> dict:
+        loss = self.model(val_batch, device=self.device)["loss"]
+
+        metrics = self.model.compute_metrics(val_batch)
+
+        return {"loss": loss, "metrics": metrics}
+
+    def log_samples(self, batch: Dict[str, Any]):
+        logging.debug("log_samples")
+        logs = self.model.log_samples(
+            batch,
+            **self.log_samples_model_kwargs,
+        )
+
+        if logs is not None:
+            N = min([logs[keys].shape[0] for keys in logs])
+        else:
+            N = 0
+
+        # Log inputs
+        if self.log_keys is not None:
+            for key in self.log_keys:
+                if key in batch:
+                    if N > 0:
+                        logs[key] = batch[key][:N]
+                    else:
+                        logs[key] = batch[key]
+
+        return logs

libcom/image_harmonization/source/src/lbm/trainer/training_config.py

@@ -0,0 +1,82 @@
+from dataclasses import field
+from typing import List, Literal, Optional, Union
+
+from pydantic.dataclasses import dataclass
+
+from ..config import BaseConfig
+
+
+@dataclass
+class TrainingConfig(BaseConfig):
+    """
+    Configuration for the training pipeline
+
+    Args:
+
+        experiment_id (str):
+            The experiment id for the training run. If not provided, a random id will be generated.
+        optimizer_name (str):
+            The optimizer to use. Default is "AdamW". Choices are "Adam", "AdamW", "Adadelta", "Adagrad", "RMSprop", "SGD"
+        optimizer_kwargs (Dict[str, Any])
+            The optimizer kwargs. Default is [{}]
+        learning_rate (float):
+            The learning rate to use. Default is 1e-3
+        lr_scheduler_name (str):
+            The learning rate scheduler to use. Default is None. Choices are "StepLR", "CosineAnnealingLR",
+            "CosineAnnealingWarmRestarts", "ReduceLROnPlateau", "ExponentialLR"
+        lr_scheduler_kwargs (Dict[str, Any])
+            The learning rate scheduler kwargs. Default is [{}]
+        lr_scheduler_interval (str):
+            The learning rate scheduler interval. Default is ["step"]. Choices are "step", "epoch"
+        lr_scheduler_frequency (int):
+            The learning rate scheduler frequency. Default is 1
+        metrics (List[str])
+            The metrics to use. Default is None
+        tracking_metrics: Optional[List[str]]
+            The metrics to track. Default is None
+        backup_every (int):
+            The frequency to backup the model. Default is 50.
+        trainable_params (Union[str, List[str]]):
+            Regexes indicateing the parameters to train.
+            Default is [["./*"]] (i.e. all parameters are trainable)
+        log_keys: Union[str, List[str]]:
+            The keys to log when sampling from the model. Default is "txt"
+        log_samples_model_kwargs (Dict[str, Any]):
+            The kwargs for logging samples from the model. Default is {
+                "max_samples": 4,
+                "num_steps": 20,
+                "input_shape": None,
+            }
+    """
+
+    experiment_id: Optional[str] = None
+    optimizer_name: Literal[
+        "Adam", "AdamW", "Adadelta", "Adagrad", "RMSprop", "SGD"
+    ] = field(default_factory=lambda: "AdamW")
+    optimizer_kwargs: Optional[dict] = field(default_factory=lambda: {})
+    learning_rate: float = field(default_factory=lambda: 1e-3)
+    lr_scheduler_name: Optional[
+        Literal[
+            "StepLR",
+            "CosineAnnealingLR",
+            "CosineAnnealingWarmRestarts",
+            "ReduceLROnPlateau",
+            "ExponentialLR",
+            None,
+        ]
+    ] = None
+    lr_scheduler_kwargs: Optional[dict] = field(default_factory=lambda: {})
+    lr_scheduler_interval: Optional[Literal["step", "epoch", None]] = "step"
+    lr_scheduler_frequency: Optional[int] = 1
+    metrics: Optional[List[str]] = None
+    tracking_metrics: Optional[List[str]] = None
+    backup_every: int = 50
+    trainable_params: List[str] = field(default_factory=lambda: ["./*"])
+    log_keys: Optional[Union[str, List[str]]] = "txt"
+    log_samples_model_kwargs: Optional[dict] = field(
+        default_factory=lambda: {
+            "max_samples": 4,
+            "num_steps": 20,
+            "input_shape": None,
+        }
+    )