quadra 0.0.1__py3-none-any.whl → 2.2.7__py3-none-any.whl

quadra/callbacks/mlflow.py

@@ -0,0 +1,291 @@
+from __future__ import annotations
+
+import glob
+import os
+from typing import Any, Literal
+
+import torch
+from pytorch_lightning import Callback, LightningModule, Trainer
+from pytorch_lightning.callbacks import LearningRateMonitor
+from pytorch_lightning.loggers import MLFlowLogger
+from pytorch_lightning.utilities import rank_zero_only
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+
+from quadra.utils.mlflow import get_mlflow_logger
+
+
+def check_minio_credentials() -> None:
+    """Check minio credentials for aws based storage such as minio.
+
+    Returns:
+        None
+    """
+    check = os.environ.get("AWS_ACCESS_KEY_ID") is not None and os.environ.get("AWS_SECRET_ACCESS_KEY") is not None
+    if not check:
+        raise ValueError(
+            "You are trying to upload mlflow artifacts, but minio credentials are not set. Please set them in your"
+            " environment variables."
+        )
+
+
+def check_file_server_dependencies() -> None:
+    """Check file dependencies as boto3.
+
+    Returns:
+        None
+    """
+    try:
+        # pylint: disable=unused-import,import-outside-toplevel
+        import boto3  # noqa
+        import minio  # noqa
+    except ImportError as e:
+        raise ImportError(
+            "You are trying to upload mlflow artifacts, but boto3 and minio are not installed. Please install them by"
+            " calling pip install minio boto3."
+        ) from e
+
+
+def validate_artifact_storage(logger: MLFlowLogger):
+    """Validate artifact storage.
+
+    Args:
+        logger: Mlflow logger from pytorch lightning.
+
+    """
+    from quadra.utils.utils import get_logger  # pylint: disable=[import-outside-toplevel]
+
+    log = get_logger(__name__)
+
+    client = logger.experiment
+    # TODO: we have to access the internal api to get the artifact uri, however there could be a better way
+    artifact_uri = client._tracking_client._get_artifact_repo(  # pylint: disable=protected-access
+        logger.run_id
+    ).artifact_uri
+    if artifact_uri.startswith("s3://"):
+        check_minio_credentials()
+        check_file_server_dependencies()
+        log.info("Mlflow artifact storage is AWS/S3 basedand credentials and dependencies are satisfied.")
+    else:
+        log.info("Mlflow artifact storage uri is %s. Validation checks are not implemented.", artifact_uri)
+
+
+class UploadCodeAsArtifact(Callback):
+    """Callback used to upload Code as artifact.
+
+    Uploads all *.py files to mlflow as an artifact, at the beginning of the run but
+        after initializing the trainer. It creates project-source folder under mlflow
+        artifacts and other necessary subfolders.
+
+    Args:
+        source_dir: Folder where all the source files are stored.
+    """
+
+    def __init__(self, source_dir: str):
+        self.source_dir = source_dir
+
+    @rank_zero_only
+    def on_test_end(self, trainer: Trainer, pl_module: LightningModule):
+        """Triggered at the end of test. Uploads all *.py files to mlflow as an artifact.
+
+        Args:
+            trainer: Pytorch Lightning trainer.
+            pl_module: Pytorch Lightning module.
+        """
+        logger = get_mlflow_logger(trainer=trainer)
+
+        if logger is None:
+            return
+
+        experiment = logger.experiment
+
+        for path in glob.glob(os.path.join(self.source_dir, "**/*.py"), recursive=True):
+            stripped_path = path.replace(self.source_dir, "")
+            if len(stripped_path.split("/")) > 1:
+                file_path_tree = "/" + "/".join(stripped_path.split("/")[:-1])
+            else:
+                file_path_tree = ""
+            experiment.log_artifact(
+                run_id=logger.run_id,
+                local_path=path,
+                artifact_path=f"project-source{file_path_tree}",
+            )
+
+
+class LogGradients(Callback):
+    """Callback used to logs of the model at the end of the of each training step.
+
+    Args:
+        norm: Norm to use for the gradient. Default is L2 norm.
+        tag: Tag to add to the gradients. If None, no tag will be added.
+        sep: Separator to use in the log.
+        round_to: Number of decimals to round the gradients to.
+        log_all_grads: If True, log all gradients, not just the total norm.
+    """
+
+    def __init__(
+        self,
+        norm: int = 2,
+        tag: str | None = None,
+        sep: str = "/",
+        round_to: int = 3,
+        log_all_grads: bool = False,
+    ):
+        self.norm = norm
+        self.tag = tag
+        self.sep = sep
+        self.round_to = round_to
+        self.log_all_grads = log_all_grads
+
+    def _grad_norm(self, named_params) -> dict:
+        """Compute the gradient norm and return it in a dictionary."""
+        grad_tag = "" if self.tag is None else "_" + self.tag
+        results = {}
+        for name, p in named_params:
+            if p.requires_grad and p.grad is not None:
+                norm = float(p.grad.data.norm(self.norm))
+                key = f"grad_norm_{self.norm}{grad_tag}{self.sep}{name}"
+                results[key] = round(norm, 3)
+        total_norm = float(torch.tensor(list(results.values())).norm(self.norm))
+        if not self.log_all_grads:
+            # clear dictionary
+            results = {}
+        results[f"grad_norm_{self.norm}_total{grad_tag}"] = round(total_norm, self.round_to)
+        return results
+
+    @rank_zero_only
+    def on_train_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+        unused: int | None = 0,
+    ) -> None:
+        """Method called at the end of the train batch
+        Args:
+            trainer: pl.trainer
+            pl_module: lightning module
+            outputs: outputs
+            batch: batch
+            batch_idx: index
+            unused: dl index.
+
+
+        Returns:
+            None
+        """
+        # pylint: disable=unused-argument
+        logger = get_mlflow_logger(trainer=trainer)
+
+        if logger is None:
+            return
+
+        named_params = pl_module.named_parameters()
+        grads = self._grad_norm(named_params)
+        logger.log_metrics(grads)
+
+
+class UploadCheckpointsAsArtifact(Callback):
+    """Callback used to upload checkpoints as artifacts.
+
+    Args:
+        ckpt_dir: Folder where all the checkpoints are stored in artifact folder.
+        ckpt_ext: Extension of checkpoint files (default: ckpt).
+        upload_best_only: Only upload best checkpoint (default: False)
+        delete_after_upload: Delete the checkpoint from local storage after uploading (default: True)
+        upload: If True, upload the checkpoints. If False, only save them on local machine.
+    """
+
+    def __init__(
+        self,
+        ckpt_dir: str = "checkpoints/",
+        ckpt_ext: str = "ckpt",
+        upload_best_only: bool = False,
+        delete_after_upload: bool = True,
+        upload: bool = True,
+    ):
+        self.ckpt_dir = ckpt_dir
+        self.upload_best_only = upload_best_only
+        self.ckpt_ext = ckpt_ext
+        self.delete_after_upload = delete_after_upload
+        self.upload = upload
+
+    @rank_zero_only
+    def on_test_end(self, trainer: Trainer, pl_module: LightningModule):
+        """Triggered at the end of test. Uploads all model checkpoints to mlflow as an artifact.
+
+        Args:
+            trainer: Pytorch Lightning trainer.
+            pl_module: Pytorch Lightning module.
+        """
+        logger = get_mlflow_logger(trainer=trainer)
+
+        if logger is None:
+            return
+
+        experiment = logger.experiment
+
+        if (
+            trainer.checkpoint_callback
+            and self.upload_best_only
+            and hasattr(trainer.checkpoint_callback, "best_model_path")
+        ):
+            if self.upload:
+                experiment.log_artifact(
+                    run_id=logger.run_id,
+                    local_path=trainer.checkpoint_callback.best_model_path,
+                    artifact_path="checkpoints",
+                )
+        else:
+            for path in glob.glob(os.path.join(self.ckpt_dir, f"**/*.{self.ckpt_ext}"), recursive=True):
+                if self.upload:
+                    experiment.log_artifact(
+                        run_id=logger.run_id,
+                        local_path=path,
+                        artifact_path="checkpoints",
+                    )
+        if self.delete_after_upload:
+            for path in glob.glob(os.path.join(self.ckpt_dir, f"**/*.{self.ckpt_ext}"), recursive=True):
+                os.remove(path)
+
+
+class LogLearningRate(LearningRateMonitor):
+    """Learning rate logger at the end of the training step/epoch.
+
+    Args:
+        logging_interval: Logging interval.
+        log_momentum: If True, log momentum as well.
+    """
+
+    def __init__(self, logging_interval: Literal["step", "epoch"] | None = None, log_momentum: bool = False):
+        super().__init__(logging_interval=logging_interval, log_momentum=log_momentum)
+
+    def on_train_batch_start(self, trainer, *args, **kwargs):
+        """Log learning rate at the beginning of the training step if logging interval is set to step."""
+        if not trainer.logger_connector.should_update_logs:
+            return
+        if self.logging_interval != "epoch":
+            logger = get_mlflow_logger(trainer=trainer)
+
+            if logger is None:
+                return
+
+            interval = "step" if self.logging_interval is None else "any"
+            latest_stat = self._extract_stats(trainer, interval)
+
+            if latest_stat:
+                logger.log_metrics(latest_stat, step=trainer.global_step)
+
+    def on_train_epoch_start(self, trainer, *args, **kwargs):
+        """Log learning rate at the beginning of the epoch if logging interval is set to epoch."""
+        if self.logging_interval != "step":
+            interval = "epoch" if self.logging_interval is None else "any"
+            latest_stat = self._extract_stats(trainer, interval)
+            logger = get_mlflow_logger(trainer=trainer)
+
+            if logger is None:
+                return
+
+            if latest_stat:
+                logger.log_metrics(latest_stat, step=trainer.global_step)

quadra/callbacks/scheduler.py

@@ -0,0 +1,69 @@
+import hydra
+import pytorch_lightning as pl
+from omegaconf import DictConfig
+from pytorch_lightning import Callback
+from pytorch_lightning.utilities import rank_zero_only
+
+from quadra.schedulers.warmup import CosineAnnealingWithLinearWarmUp
+from quadra.utils.utils import get_logger
+
+log = get_logger(__name__)
+
+
+class WarmupInit(Callback):
+    """Custom callback used to setup a warmup scheduler.
+
+    Args:
+        scheduler_config: scheduler configuration.
+    """
+
+    def __init__(
+        self,
+        scheduler_config: DictConfig,
+    ) -> None:
+        self.scheduler_config = scheduler_config
+
+    @rank_zero_only
+    def on_fit_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        """Called when fit begins."""
+        if not hasattr(trainer, "datamodule"):
+            raise ValueError("Trainer must have a datamodule attribute.")
+
+        if not any(isinstance(s.scheduler, CosineAnnealingWithLinearWarmUp) for s in trainer.lr_scheduler_configs):
+            return
+
+        log.info("Using warmup scheduler, forcing optimizer learning rate to zero.")
+        for i, _ in enumerate(trainer.optimizers):
+            for param_group in trainer.optimizers[i].param_groups:
+                param_group["lr"] = 0.0
+            trainer.optimizers[i].defaults["lr"] = 0.0
+
+        batch_size = trainer.datamodule.batch_size
+        train_dataloader = trainer.datamodule.train_dataloader()
+        len_train_dataloader = len(train_dataloader)
+        if isinstance(trainer.device_ids, list) and pl_module.device.type == "cuda":
+            num_gpus = len(trainer.device_ids)
+            len_train_dataloader = len_train_dataloader // num_gpus
+            if not train_dataloader.drop_last:
+                len_train_dataloader += int((len_train_dataloader % num_gpus) != 0)
+
+        if len_train_dataloader == 1:
+            log.warning(
+                "From this dataset size, we can only generate single batch. The batch size will be set as lenght of"
+                " the dataset "
+            )
+            batch_size = len(train_dataloader.dataset)
+
+        if isinstance(trainer.device_ids, list) and pl_module.device.type == "cuda":
+            batch_size = batch_size * len(trainer.device_ids)
+
+        scheduler = hydra.utils.instantiate(
+            self.scheduler_config,
+            optimizer=pl_module.optimizer,
+            batch_size=batch_size,
+            len_loader=len_train_dataloader,
+        )
+
+        for i, s in enumerate(trainer.lr_scheduler_configs):
+            if isinstance(s.scheduler, CosineAnnealingWithLinearWarmUp):
+                trainer.lr_scheduler_configs[i].scheduler = scheduler

quadra/configs/backbone/caformer_m36.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: caformer_m36.sail_in22k_ft_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 576

quadra/configs/backbone/caformer_s36.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: caformer_s36.sail_in22k_ft_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 512

quadra/configs/backbone/convnextv2_base.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: convnextv2_base.fcmae
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 1024

quadra/configs/backbone/convnextv2_femto.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: convnextv2_femto.fcmae
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 384

quadra/configs/backbone/convnextv2_tiny.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: convnextv2_tiny.fcmae
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 768

quadra/configs/backbone/dino_vitb8.yaml

@@ -0,0 +1,12 @@
+model:
+  _target_: quadra.models.classification.TorchHubNetworkBuilder
+  repo_or_dir: facebookresearch/dino:main
+  model_name: dino_vitb8
+  pretrained: true
+  freeze: false
+  hyperspherical: false
+metadata:
+  input_size: 224
+  output_dim: 768
+  patch_size: 8
+  nb_heads: 12

quadra/configs/backbone/dino_vits8.yaml

@@ -0,0 +1,12 @@
+model:
+  _target_: quadra.models.classification.TorchHubNetworkBuilder
+  repo_or_dir: facebookresearch/dino:main
+  model_name: dino_vits8
+  pretrained: true
+  freeze: false
+  hyperspherical: false
+metadata:
+  input_size: 224
+  output_dim: 384
+  patch_size: 8
+  nb_heads: 6

quadra/configs/backbone/dinov2_vitb14.yaml

@@ -0,0 +1,12 @@
+model:
+  _target_: quadra.models.classification.TorchHubNetworkBuilder
+  repo_or_dir: facebookresearch/dinov2
+  model_name: dinov2_vitb14
+  pretrained: true
+  freeze: false
+  hyperspherical: false
+metadata:
+  input_size: 224
+  output_dim: 768
+  patch_size: 14
+  nb_heads: 12

quadra/configs/backbone/dinov2_vits14.yaml

@@ -0,0 +1,12 @@
+model:
+  _target_: quadra.models.classification.TorchHubNetworkBuilder
+  repo_or_dir: facebookresearch/dinov2
+  model_name: dinov2_vits14
+  pretrained: true
+  freeze: false
+  hyperspherical: false
+metadata:
+  input_size: 224
+  output_dim: 384
+  patch_size: 14
+  nb_heads: 6

quadra/configs/backbone/efficientnet_b0.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: tf_efficientnet_b0.ns_jft_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 1280

quadra/configs/backbone/efficientnet_b1.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: tf_efficientnet_b1.ns_jft_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 1280

quadra/configs/backbone/efficientnet_b2.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: tf_efficientnet_b2.ns_jft_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 260
+  output_dim: 1408

quadra/configs/backbone/efficientnet_b3.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: tf_efficientnet_b3.ns_jft_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 300
+  output_dim: 1536

quadra/configs/backbone/efficientnetv2_s.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: tf_efficientnetv2_s_in21ft1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 384
+  output_dim: 1280

quadra/configs/backbone/levit_128s.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: levit_128s
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 384

quadra/configs/backbone/mnasnet0_5.yaml

@@ -0,0 +1,9 @@
+model:
+  _target_: quadra.models.classification.TorchVisionNetworkBuilder
+  model_name: mnasnet0_5
+  pretrained: true
+  freeze: false
+  hyperspherical: false
+metadata:
+  input_size: 224
+  output_dim: 1280

quadra/configs/backbone/resnet101.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: resnet101
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 2048

quadra/configs/backbone/resnet18.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: resnet18.tv_in1k # Use torchvision weights
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 512

quadra/configs/backbone/resnet18_ssl.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: ssl_resnet18
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 512

quadra/configs/backbone/resnet50.yaml

@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: resnet50
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 2048

quadra/configs/backbone/smp.yaml

@@ -0,0 +1,9 @@
+model:
+  _target_: quadra.modules.backbone.create_smp_backbone
+  arch: unet
+  encoder_name: resnet18
+  encoder_weights: imagenet
+  freeze_encoder: True
+  in_channels: 3
+  num_classes: 1
+  activation: null

quadra/configs/backbone/tiny_vit_21m_224.yaml

@@ -0,0 +1,9 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: tiny_vit_21m_224.dist_in22k_ft_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 576
+  num_heads: 18 # Is it correct?

quadra/configs/backbone/unetr.yaml

@@ -0,0 +1,15 @@
+model:
+  _target_: monai.networks.nets.unetr.UNETR
+  in_channels: 3
+  out_channels: 1
+  img_size: [448, 448]
+  feature_size: 16
+  hidden_size: 384 # 192
+  mlp_dim: 1536 # 768
+  num_heads: 8 # 3
+  pos_embed: conv
+  norm_name: instance
+  conv_block: true
+  res_block: true
+  dropout_rate: 0
+  spatial_dims: 2

quadra/configs/backbone/vit16_base.yaml

@@ -0,0 +1,9 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: vit_base_patch16_224
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 768
+  nb_heads: 6

quadra/configs/backbone/vit16_small.yaml

@@ -0,0 +1,9 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: vit_small_patch16_224
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 384
+  nb_heads: 6

quadra/configs/backbone/vit16_tiny.yaml

@@ -0,0 +1,9 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: vit_tiny_patch16_224
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 192
+  nb_heads: 3

quadra/configs/backbone/xcit_tiny_24_p8_224.yaml

@@ -0,0 +1,9 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: xcit_tiny_24_p8_224
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 192
+  nb_heads: 4