Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# Copyright (c) Facebook, Inc. and its affiliates.2import logging3import math4from bisect import bisect_right5from typing import List6import torch7from fvcore.common.param_scheduler import (8    CompositeParamScheduler,9    ConstantParamScheduler,10    LinearParamScheduler,11    ParamScheduler,12)13 14try:15    from torch.optim.lr_scheduler import LRScheduler16except ImportError:17    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler18 19logger = logging.getLogger(__name__)20 21 22class WarmupParamScheduler(CompositeParamScheduler):23    """24    Add an initial warmup stage to another scheduler.25    """26 27    def __init__(28        self,29        scheduler: ParamScheduler,30        warmup_factor: float,31        warmup_length: float,32        warmup_method: str = "linear",33        rescale_interval: bool = False,34    ):35        """36        Args:37            scheduler: warmup will be added at the beginning of this scheduler38            warmup_factor: the factor w.r.t the initial value of ``scheduler``, e.g. 0.00139            warmup_length: the relative length (in [0, 1]) of warmup steps w.r.t the entire40                training, e.g. 0.0141            warmup_method: one of "linear" or "constant"42            rescale_interval: whether we will rescale the interval of the scheduler after43                warmup44        """45        # the value to reach when warmup ends46        end_value = scheduler(0.0) if rescale_interval else scheduler(warmup_length)47        start_value = warmup_factor * scheduler(0.0)48        if warmup_method == "constant":49            warmup = ConstantParamScheduler(start_value)50        elif warmup_method == "linear":51            warmup = LinearParamScheduler(start_value, end_value)52        else:53            raise ValueError("Unknown warmup method: {}".format(warmup_method))54        super().__init__(55            [warmup, scheduler],56            interval_scaling=["rescaled", "rescaled" if rescale_interval else "fixed"],57            lengths=[warmup_length, 1 - warmup_length],58        )59 60 61class LRMultiplier(LRScheduler):62    """63    A LRScheduler which uses fvcore :class:`ParamScheduler` to multiply the64    learning rate of each param in the optimizer.65    Every step, the learning rate of each parameter becomes its initial value66    multiplied by the output of the given :class:`ParamScheduler`.67 68    The absolute learning rate value of each parameter can be different.69    This scheduler can be used as long as the relative scale among them do70    not change during training.71 72    Examples:73    ::74        LRMultiplier(75            opt,76            WarmupParamScheduler(77                MultiStepParamScheduler(78                    [1, 0.1, 0.01],79                    milestones=[60000, 80000],80                    num_updates=90000,81                ), 0.001, 100 / 9000082            ),83            max_iter=9000084        )85    """86 87    # NOTES: in the most general case, every LR can use its own scheduler.88    # Supporting this requires interaction with the optimizer when its parameter89    # group is initialized. For example, classyvision implements its own optimizer90    # that allows different schedulers for every parameter group.91    # To avoid this complexity, we use this class to support the most common cases92    # where the relative scale among all LRs stay unchanged during training.  In this93    # case we only need a total of one scheduler that defines the relative LR multiplier.94 95    def __init__(96        self,97        optimizer: torch.optim.Optimizer,98        multiplier: ParamScheduler,99        max_iter: int,100        last_iter: int = -1,101    ):102        """103        Args:104            optimizer, last_iter: See ``torch.optim.lr_scheduler.LRScheduler``.105                ``last_iter`` is the same as ``last_epoch``.106            multiplier: a fvcore ParamScheduler that defines the multiplier on107                every LR of the optimizer108            max_iter: the total number of training iterations109        """110        if not isinstance(multiplier, ParamScheduler):111            raise ValueError(112                "_LRMultiplier(multiplier=) must be an instance of fvcore "113                f"ParamScheduler. Got {multiplier} instead."114            )115        self._multiplier = multiplier116        self._max_iter = max_iter117        super().__init__(optimizer, last_epoch=last_iter)118 119    def state_dict(self):120        # fvcore schedulers are stateless. Only keep pytorch scheduler states121        return {"base_lrs": self.base_lrs, "last_epoch": self.last_epoch}122 123    def get_lr(self) -> List[float]:124        multiplier = self._multiplier(self.last_epoch / self._max_iter)125        return [base_lr * multiplier for base_lr in self.base_lrs]126 127 128"""129Content below is no longer needed!130"""131 132# NOTE: PyTorch's LR scheduler interface uses names that assume the LR changes133# only on epoch boundaries. We typically use iteration based schedules instead.134# As a result, "epoch" (e.g., as in self.last_epoch) should be understood to mean135# "iteration" instead.136 137# FIXME: ideally this would be achieved with a CombinedLRScheduler, separating138# MultiStepLR with WarmupLR but the current LRScheduler design doesn't allow it.139 140 141class WarmupMultiStepLR(LRScheduler):142    def __init__(143        self,144        optimizer: torch.optim.Optimizer,145        milestones: List[int],146        gamma: float = 0.1,147        warmup_factor: float = 0.001,148        warmup_iters: int = 1000,149        warmup_method: str = "linear",150        last_epoch: int = -1,151    ):152        logger.warning(153            "WarmupMultiStepLR is deprecated! Use LRMultipilier with fvcore ParamScheduler instead!"154        )155        if not list(milestones) == sorted(milestones):156            raise ValueError(157                "Milestones should be a list of" " increasing integers. Got {}", milestones158            )159        self.milestones = milestones160        self.gamma = gamma161        self.warmup_factor = warmup_factor162        self.warmup_iters = warmup_iters163        self.warmup_method = warmup_method164        super().__init__(optimizer, last_epoch)165 166    def get_lr(self) -> List[float]:167        warmup_factor = _get_warmup_factor_at_iter(168            self.warmup_method, self.last_epoch, self.warmup_iters, self.warmup_factor169        )170        return [171            base_lr * warmup_factor * self.gamma ** bisect_right(self.milestones, self.last_epoch)172            for base_lr in self.base_lrs173        ]174 175    def _compute_values(self) -> List[float]:176        # The new interface177        return self.get_lr()178 179 180class WarmupCosineLR(LRScheduler):181    def __init__(182        self,183        optimizer: torch.optim.Optimizer,184        max_iters: int,185        warmup_factor: float = 0.001,186        warmup_iters: int = 1000,187        warmup_method: str = "linear",188        last_epoch: int = -1,189    ):190        logger.warning(191            "WarmupCosineLR is deprecated! Use LRMultipilier with fvcore ParamScheduler instead!"192        )193        self.max_iters = max_iters194        self.warmup_factor = warmup_factor195        self.warmup_iters = warmup_iters196        self.warmup_method = warmup_method197        super().__init__(optimizer, last_epoch)198 199    def get_lr(self) -> List[float]:200        warmup_factor = _get_warmup_factor_at_iter(201            self.warmup_method, self.last_epoch, self.warmup_iters, self.warmup_factor202        )203        # Different definitions of half-cosine with warmup are possible. For204        # simplicity we multiply the standard half-cosine schedule by the warmup205        # factor. An alternative is to start the period of the cosine at warmup_iters206        # instead of at 0. In the case that warmup_iters << max_iters the two are207        # very close to each other.208        return [209            base_lr210            * warmup_factor211            * 0.5212            * (1.0 + math.cos(math.pi * self.last_epoch / self.max_iters))213            for base_lr in self.base_lrs214        ]215 216    def _compute_values(self) -> List[float]:217        # The new interface218        return self.get_lr()219 220 221def _get_warmup_factor_at_iter(222    method: str, iter: int, warmup_iters: int, warmup_factor: float223) -> float:224    """225    Return the learning rate warmup factor at a specific iteration.226    See :paper:`ImageNet in 1h` for more details.227 228    Args:229        method (str): warmup method; either "constant" or "linear".230        iter (int): iteration at which to calculate the warmup factor.231        warmup_iters (int): the number of warmup iterations.232        warmup_factor (float): the base warmup factor (the meaning changes according233            to the method used).234 235    Returns:236        float: the effective warmup factor at the given iteration.237    """238    if iter >= warmup_iters:239        return 1.0240 241    if method == "constant":242        return warmup_factor243    elif method == "linear":244        alpha = iter / warmup_iters245        return warmup_factor * (1 - alpha) + alpha246    else:247        raise ValueError("Unknown warmup method: {}".format(method))248