plato-learn 1.1__py3-none-any.whl

plato/trainers/loss_criterion.py

@@ -0,0 +1,70 @@
+"""
+Obtaining the loss criterion for training workloads according to the configuration file.
+"""
+
+from typing import Union
+
+from torch import nn
+from lightly import loss
+
+from plato.config import Config
+
+
+def get(**kwargs: Union[str, dict]):
+    """Get a loss function with its name from the configuration file."""
+    registered_loss_criterion = {
+        "L1Loss": nn.L1Loss,
+        "MSELoss": nn.MSELoss,
+        "BCELoss": nn.BCELoss,
+        "BCEWithLogitsLoss": nn.BCEWithLogitsLoss,
+        "NLLLoss": nn.NLLLoss,
+        "PoissonNLLLoss": nn.PoissonNLLLoss,
+        "CrossEntropyLoss": nn.CrossEntropyLoss,
+        "HingeEmbeddingLoss": nn.HingeEmbeddingLoss,
+        "MarginRankingLoss": nn.MarginRankingLoss,
+        "TripletMarginLoss": nn.TripletMarginLoss,
+        "KLDivLoss": nn.KLDivLoss,
+    }
+
+    ssl_loss_criterion = {
+        "NegativeCosineSimilarity": loss.NegativeCosineSimilarity,
+        "NTXentLoss": loss.NTXentLoss,
+        "BarlowTwinsLoss": loss.BarlowTwinsLoss,
+        "DCLLoss": loss.DCLLoss,
+        "DCLWLoss": loss.DCLWLoss,
+        "DINOLoss": loss.DINOLoss,
+        "PMSNCustomLoss": loss.PMSNCustomLoss,
+        "SwaVLoss": loss.SwaVLoss,
+        "PMSNLoss": loss.PMSNLoss,
+        "SymNegCosineSimilarityLoss": loss.SymNegCosineSimilarityLoss,
+        "TiCoLoss": loss.TiCoLoss,
+        "VICRegLoss": loss.VICRegLoss,
+        "VICRegLLoss": loss.VICRegLLoss,
+        "MSNLoss": loss.MSNLoss,
+    }
+
+    registered_loss_criterion.update(ssl_loss_criterion)
+
+    loss_criterion_name = (
+        kwargs["loss_criterion"]
+        if "loss_criterion" in kwargs
+        else (
+            Config().trainer.loss_criterion
+            if hasattr(Config.trainer, "loss_criterion")
+            else "CrossEntropyLoss"
+        )
+    )
+
+    loss_criterion_params = (
+        kwargs["loss_criterion_params"]
+        if "loss_criterion_params" in kwargs
+        else (
+            Config().parameters.loss_criterion._asdict()
+            if hasattr(Config.parameters, "loss_criterion")
+            else {}
+        )
+    )
+
+    loss_criterion = registered_loss_criterion.get(loss_criterion_name)
+
+    return loss_criterion(**loss_criterion_params)

plato/trainers/lr_schedulers.py

@@ -0,0 +1,252 @@
+"""
+Returns a learning rate scheduler according to the configuration.
+"""
+
+import bisect
+import sys
+from types import SimpleNamespace
+from typing import Union
+
+import numpy as np
+from timm import scheduler
+from torch import optim
+
+from plato.config import Config
+
+
+def get(
+    optimizer: optim.Optimizer, iterations_per_epoch: int, **kwargs: Union[str, dict]
+):
+    """Returns a learning rate scheduler according to the configuration."""
+
+    registered_schedulers = {
+        "CosineAnnealingLR": optim.lr_scheduler.CosineAnnealingLR,
+        "LambdaLR": optim.lr_scheduler.LambdaLR,
+        "MultiStepLR": optim.lr_scheduler.MultiStepLR,
+        "StepLR": optim.lr_scheduler.StepLR,
+        "ReduceLROnPlateau": optim.lr_scheduler.ReduceLROnPlateau,
+        "ConstantLR": optim.lr_scheduler.ConstantLR,
+        "LinearLR": optim.lr_scheduler.LinearLR,
+        "ExponentialLR": optim.lr_scheduler.ExponentialLR,
+        "CyclicLR": optim.lr_scheduler.CyclicLR,
+        "CosineAnnealingWarmRestarts": optim.lr_scheduler.CosineAnnealingWarmRestarts,
+    }
+
+    registered_factories = {
+        "timm": scheduler.create_scheduler,
+    }
+
+    _scheduler = (
+        kwargs["lr_scheduler"]
+        if "lr_scheduler" in kwargs
+        else Config().trainer.lr_scheduler
+    )
+    lr_params = (
+        kwargs["lr_params"]
+        if "lr_params" in kwargs
+        else Config().parameters.learning_rate._asdict()
+    )
+
+    # First, look up the registered factories of LR schedulers
+    if _scheduler in registered_factories:
+        scheduler_args = SimpleNamespace(**lr_params)
+        scheduler_args.epochs = Config().trainer.epochs
+        lr_scheduler, __ = registered_factories[_scheduler](
+            args=scheduler_args, optimizer=optimizer
+        )
+        return lr_scheduler
+
+    # The list containing the learning rate schedulers that must be returned or
+    # the learning rate schedulers that ChainedScheduler or SequentialLR will
+    # take as an argument.
+    returned_schedulers = []
+
+    use_chained = False
+    use_sequential = False
+    if "ChainedScheduler" in _scheduler:
+        use_chained = True
+        lr_scheduler = [
+            sched for sched in _scheduler.split(",") if sched != ("ChainedScheduler")
+        ]
+    elif "SequentialLR" in _scheduler:
+        use_sequential = True
+        lr_scheduler = [
+            sched for sched in _scheduler.split(",") if sched != ("SequentialLR")
+        ]
+    else:
+        lr_scheduler = [_scheduler]
+
+    for _scheduler in lr_scheduler:
+        retrieved_scheduler = registered_schedulers.get(_scheduler)
+
+        if retrieved_scheduler is None:
+            sys.exit("Error: Unknown learning rate scheduler.")
+
+        if _scheduler == "CosineAnnealingLR":
+            returned_schedulers.append(
+                retrieved_scheduler(
+                    optimizer, iterations_per_epoch * Config().trainer.epochs
+                )
+            )
+        elif _scheduler == "LambdaLR":
+            lambdas = [lambda it: 1.0]
+
+            if "gamma" in lr_params and "milestone_steps" in lr_params:
+                milestones = [
+                    Step.from_str(x, iterations_per_epoch).iteration
+                    for x in lr_params["milestone_steps"].split(",")
+                ]
+                lambdas.append(
+                    lambda it, milestones=milestones: lr_params["gamma"]
+                    ** bisect.bisect(milestones, it)
+                )
+
+            # Add a linear learning rate warmup if specified
+            if "warmup_steps" in lr_params:
+                warmup_iters = Step.from_str(
+                    lr_params["warmup_steps"], iterations_per_epoch
+                ).iteration
+                lambdas.append(
+                    lambda it, warmup_iters=warmup_iters: min(1.0, it / warmup_iters)
+                )
+            returned_schedulers.append(
+                retrieved_scheduler(
+                    optimizer,
+                    lambda it, lambdas=lambdas: np.prod([l(it) for l in lambdas]),
+                )
+            )
+        elif _scheduler == "MultiStepLR":
+            milestones = [
+                int(x.split("ep")[0]) for x in lr_params["milestone_steps"].split(",")
+            ]
+            returned_schedulers.append(
+                retrieved_scheduler(
+                    optimizer, milestones=milestones, gamma=lr_params["gamma"]
+                )
+            )
+        else:
+            returned_schedulers.append(retrieved_scheduler(optimizer, **lr_params))
+
+    if use_chained:
+        return optim.lr_scheduler.ChainedScheduler(returned_schedulers)
+
+    if use_sequential:
+        sequential_milestones = (
+            Config().trainer.lr_sequential_milestones
+            if hasattr(Config().trainer, "lr_sequential_milestones")
+            else 2
+        )
+        sequential_milestones = [
+            int(epoch) for epoch in sequential_milestones.split(",")
+        ]
+
+        return optim.lr_scheduler.SequentialLR(
+            optimizer, returned_schedulers, sequential_milestones
+        )
+    else:
+        return returned_schedulers[0]
+
+
+class Step:
+    """Represents a particular step of training."""
+
+    def __init__(self, iteration: int, iterations_per_epoch: int) -> None:
+        if iteration < 0:
+            raise ValueError("iteration must >= 0.")
+        if iterations_per_epoch <= 0:
+            raise ValueError("iterations_per_epoch must be > 0.")
+        self._iteration = iteration
+        self.iterations_per_epoch = iterations_per_epoch
+
+    @staticmethod
+    def str_is_zero(s: str) -> bool:
+        return s in ["0ep", "0it", "0ep0it"]
+
+    @staticmethod
+    def from_iteration(iteration: int, iterations_per_epoch: int) -> "Step":
+        return Step(iteration, iterations_per_epoch)
+
+    @staticmethod
+    def from_epoch(epoch: int, iteration: int, iterations_per_epoch: int) -> "Step":
+        return Step(epoch * iterations_per_epoch + iteration, iterations_per_epoch)
+
+    @staticmethod
+    def from_str(s: str, iterations_per_epoch: int) -> "Step":
+        """Creates a step from a string that describes the number of epochs, iterations, or both.
+
+        Epochs: '120ep'
+        Iterations: '2000it'
+        Both: '120ep50it'"""
+
+        if "ep" in s and "it" in s:
+            ep = int(s.split("ep")[0])
+            it = int(s.split("ep")[1].split("it")[0])
+            if s != "{}ep{}it".format(ep, it):
+                raise ValueError(f"Malformed string step: {s}")
+            return Step.from_epoch(ep, it, iterations_per_epoch)
+        elif "ep" in s:
+            ep = int(s.split("ep")[0])
+            if s != "{}ep".format(ep):
+                raise ValueError(f"Malformed string step: {s}")
+            return Step.from_epoch(ep, 0, iterations_per_epoch)
+        elif "it" in s:
+            it = int(s.split("it")[0])
+            if s != "{}it".format(it):
+                raise ValueError(f"Malformed string step: {s}")
+            return Step.from_iteration(it, iterations_per_epoch)
+        else:
+            raise ValueError(f"Malformed string step: {s}")
+
+    @staticmethod
+    def zero(iterations_per_epoch: int) -> "Step":
+        return Step(0, iterations_per_epoch)
+
+    @property
+    def iteration(self):
+        """The overall number of steps of training completed so far."""
+        return self._iteration
+
+    @property
+    def ep(self):
+        """The current epoch of training."""
+        return self._iteration // self.iterations_per_epoch
+
+    @property
+    def it(self):
+        """The iteration within the current epoch of training."""
+        return self._iteration % self.iterations_per_epoch
+
+    def _check(self, other):
+        if not isinstance(other, Step):
+            raise ValueError(f"Invalid type for other: {other}.")
+        if self.iterations_per_epoch != other.iterations_per_epoch:
+            raise ValueError(
+                "Cannot compare steps when epochs are of different lengths."
+            )
+
+    def __lt__(self, other):
+        self._check(other)
+        return self._iteration < other.iteration
+
+    def __le__(self, other):
+        self._check(other)
+        return self._iteration <= other.iteration
+
+    def __eq__(self, other):
+        self._check(other)
+        return self._iteration == other.iteration
+
+    def __ne__(self, other):
+        self._check(other)
+        return self._iteration != other.iteration
+
+    def __gt__(self, other):
+        self._check(other)
+        return self._iteration > other.iteration
+
+    def __ge__(self, other):
+        self._check(other)
+        return self._iteration >= other.iteration
+
+    def __str__(self):
+        return f"(Iteration {self._iteration}; Iterations per Epoch: {self.iterations_per_epoch})"

plato/trainers/optimizers.py

@@ -0,0 +1,53 @@
+"""
+Optimizers for training workloads.
+"""
+
+from typing import Union
+
+import torch_optimizer as torch_optim
+from torch import optim
+from timm import optim as timm_optim
+
+from plato.config import Config
+
+
+def get(model, **kwargs: Union[str, dict]) -> optim.Optimizer:
+    """Get an optimizer with its name and parameters obtained from the configuration file."""
+    registered_optimizers = {
+        "Adam": optim.Adam,
+        "Adadelta": optim.Adadelta,
+        "Adagrad": optim.Adagrad,
+        "AdaHessian": torch_optim.Adahessian,
+        "AdamW": optim.AdamW,
+        "SparseAdam": optim.SparseAdam,
+        "Adamax": optim.Adamax,
+        "ASGD": optim.ASGD,
+        "LBFGS": optim.LBFGS,
+        "NAdam": optim.NAdam,
+        "RAdam": optim.RAdam,
+        "RMSprop": optim.RMSprop,
+        "Rprop": optim.Rprop,
+        "SGD": optim.SGD,
+        "LARS": timm_optim.lars.Lars,
+    }
+
+    optimizer_name = (
+        kwargs["optimizer_name"]
+        if "optimizer_name" in kwargs
+        else Config().trainer.optimizer
+    )
+    optimizer_params = (
+        kwargs["optimizer_params"]
+        if "optimizer_params" in kwargs
+        else Config().parameters.optimizer._asdict()
+    )
+
+    # Ensure eps is a float
+    if "eps" in optimizer_params:
+        optimizer_params["eps"] = float(optimizer_params["eps"])
+
+    optimizer = registered_optimizers.get(optimizer_name)
+    if optimizer is not None:
+        return optimizer(model.parameters(), **optimizer_params)
+
+    raise ValueError(f"No such optimizer: {optimizer_name}")

plato/trainers/pascal_voc.py

@@ -0,0 +1,80 @@
+"""
+A customized trainer for image segmentation on PASCAL VOC dataset (2012).
+"""
+
+import torch.nn as nn
+import torch
+import numpy as np
+
+from plato.trainers import basic
+
+
+class Evaluator(object):
+    def __init__(self, num_class):
+        self.num_class = num_class
+        self.confusion_matrix = np.zeros((self.num_class,) * 2)
+
+    def Mean_Intersection_over_Union(self):
+        MIoU = np.diag(self.confusion_matrix) / (
+            np.sum(self.confusion_matrix, axis=1)
+            + np.sum(self.confusion_matrix, axis=0)
+            - np.diag(self.confusion_matrix)
+        )
+        MIoU = np.nanmean(MIoU)
+        return MIoU
+
+    def _generate_matrix(self, gt_image, pre_image):
+        mask = (gt_image >= 0) & (gt_image < self.num_class)
+        label = self.num_class * gt_image[mask].astype("int") + pre_image[mask]
+        count = np.bincount(label, minlength=self.num_class**2)
+        confusion_matrix = count.reshape(self.num_class, self.num_class)
+        return confusion_matrix
+
+    def add_batch(self, gt_image, pre_image):
+        assert gt_image.shape == pre_image.shape
+        self.confusion_matrix += self._generate_matrix(gt_image, pre_image)
+
+    def reset(self):
+        self.confusion_matrix = np.zeros((self.num_class,) * 2)
+
+
+class Trainer(basic.Trainer):
+    """The federated learning trainer for the image segmentation on PASCAL VOC"""
+
+    def __init__(self, model=None, **kwargs):
+        """Initializing the trainer with the provided model.
+
+        Arguments:
+        model: The model to train.
+        client_id: The ID of the client using this trainer (optional).
+        """
+        super().__init__(model)
+
+        self.loss_criterion = nn.BCEWithLogitsLoss()
+        self.num_class = 20
+
+    def test_model(self, config, testset, sampler=None, **kwargs):
+        test_loader = torch.utils.data.DataLoader(
+            testset, batch_size=config["batch_size"], shuffle=False
+        )
+
+        total = 0
+        evaluator = Evaluator(self.num_class)
+        evaluator.reset()
+        with torch.no_grad():
+            for examples, labels in test_loader:
+                examples, labels = examples.to(self.device), labels.to(self.device)
+
+                outputs = self.model(examples)
+
+                _, predicted = torch.max(outputs.data, 1)
+                total += labels.size(0)
+                labels = torch.squeeze(labels, 1).cpu().numpy()
+                predicted = predicted.cpu().numpy()
+                print("shape of pred: ", predicted.shape)
+                print("shape of labels: ", labels.shape)
+                evaluator.add_batch(labels, predicted)
+
+            accuracy = evaluator.Mean_Intersection_over_Union()
+
+        return accuracy

plato/trainers/registry.py

@@ -0,0 +1,44 @@
+"""
+Having a registry of all available classes is convenient for retrieving an instance
+based on a configuration at run-time.
+"""
+
+import logging
+
+from plato.config import Config
+
+
+from plato.trainers import basic, diff_privacy, pascal_voc, gan, split_learning
+
+registered_trainers = {
+    "basic": basic.Trainer,
+    "timm_basic": basic.TrainerWithTimmScheduler,
+    "diff_privacy": diff_privacy.Trainer,
+    "pascal_voc": pascal_voc.Trainer,
+    "gan": gan.Trainer,
+    "split_learning": split_learning.Trainer,
+}
+
+
+def get(model=None, callbacks=None):
+    """Get the trainer with the provided name."""
+    trainer_name = Config().trainer.type
+    logging.info("Trainer: %s", trainer_name)
+
+    if Config().trainer.model_name == "yolov8":
+        from plato.trainers import yolov8
+
+        return yolov8.Trainer()
+    elif Config().trainer.type == "HuggingFace":
+        from plato.trainers import huggingface
+
+        return huggingface.Trainer(model=model, callbacks=callbacks)
+
+    elif Config().trainer.type == "self_supervised_learning":
+        from plato.trainers import self_supervised_learning
+
+        return self_supervised_learning.Trainer(model=model, callbacks=callbacks)
+    elif trainer_name in registered_trainers:
+        return registered_trainers[trainer_name](model=model, callbacks=callbacks)
+    else:
+        raise ValueError(f"No such trainer: {trainer_name}")