From d3d0edf1f728448bfa989969239b3f5190639b8d Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 15 Aug 2026 09:41:59 -0300 Subject: [PATCH] refactor(optim): move LARS and the schedulers into lightly.optim --- lightly/optim/__init__.py | 19 +++ lightly/optim/lars.py | 166 ++++++++++++++++++++ lightly/optim/param_groups.py | 38 +++++ lightly/optim/schedulers.py | 237 +++++++++++++++++++++++++++++ lightly/utils/lars.py | 167 +------------------- lightly/utils/scheduler.py | 252 ++----------------------------- tests/optim/__init__.py | 0 tests/optim/test_param_groups.py | 62 ++++++++ 8 files changed, 540 insertions(+), 401 deletions(-) create mode 100644 lightly/optim/__init__.py create mode 100644 lightly/optim/lars.py create mode 100644 lightly/optim/param_groups.py create mode 100644 lightly/optim/schedulers.py create mode 100644 tests/optim/__init__.py create mode 100644 tests/optim/test_param_groups.py diff --git a/lightly/optim/__init__.py b/lightly/optim/__init__.py new file mode 100644 index 000000000..51e7bf4e4 --- /dev/null +++ b/lightly/optim/__init__.py @@ -0,0 +1,19 @@ +"""Optimisers, learning-rate schedules and parameter groups.""" + +from lightly.optim.lars import LARS +from lightly.optim.param_groups import param_groups +from lightly.optim.schedulers import ( + CosineWarmupScheduler, + cosine_schedule, + cosine_warmup_schedule, + linear_warmup_schedule, +) + +__all__ = [ + "CosineWarmupScheduler", + "LARS", + "cosine_schedule", + "cosine_warmup_schedule", + "linear_warmup_schedule", + "param_groups", +] diff --git a/lightly/optim/lars.py b/lightly/optim/lars.py new file mode 100644 index 000000000..44bf43e7b --- /dev/null +++ b/lightly/optim/lars.py @@ -0,0 +1,166 @@ +from typing import Any, Callable, Dict, Optional, overload + +import torch +from torch.optim.optimizer import Optimizer + + +class LARS(Optimizer): + """Extends SGD in PyTorch with LARS scaling from the paper "Large batch training of + Convolutional Networks" [0]. + + Implementation from PyTorch Lightning Bolts [1]. + + - [0]: https://arxiv.org/pdf/1708.03888.pdf + - [1]: https://github.com/Lightning-Universe/lightning-bolts/blob/2dfe45a4cf050f120d10981c45cfa2c785a1d5e6/pl_bolts/optimizers/lars.py#L1 + + Args: + params: + Iterable of parameters to optimize or dicts defining parameter groups. + lr: + Learning rate + momentum: + Momentum factor. + weight_decay: + Weight decay (L2 penalty). + dampening: + Dampening for momentum. + nesterov: + Enables Nesterov momentum. + trust_coefficient: + Trust coefficient for computing learning rate. + eps: + Eps for division denominator. + + Example: + >>> model = torch.nn.Linear(10, 1) + >>> input = torch.Tensor(10) + >>> target = torch.Tensor([1.]) + >>> loss_fn = lambda input, target: (input - target) ** 2 + >>> optimizer = LARS(model.parameters(), lr=0.1, momentum=0.9) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + .. note:: + The application of momentum in the SGD part is modified according to + the PyTorch standards. LARS scaling fits into the equation in the + following fashion. + + .. math:: + \begin{aligned} + g_{t+1} & = \text{lars_lr} * (\beta * p_{t} + g_{t+1}), \\ + v_{t+1} & = \\mu * v_{t} + g_{t+1}, \\ + p_{t+1} & = p_{t} - \text{lr} * v_{t+1}, + \\end{aligned} + + where :math:`p`, :math:`g`, :math:`v`, :math:`\\mu` and :math:`\beta` denote the + parameters, gradient, velocity, momentum, and weight decay respectively. + The :math:`lars_lr` is defined by Eq. 6 in the paper. + The Nesterov version is analogously modified. + + .. warning:: + Parameters with weight decay set to 0 will automatically be excluded from + layer-wise LR scaling. This is to ensure consistency with papers like SimCLR + and BYOL. + """ + + def __init__( + self, + params: Any, + lr: float, + momentum: float = 0, + dampening: float = 0, + weight_decay: float = 0, + nesterov: bool = False, + trust_coefficient: float = 0.001, + eps: float = 1e-8, + ): + if lr < 0.0: + raise ValueError(f"Invalid learning rate: {lr}") + if momentum < 0.0: + raise ValueError(f"Invalid momentum value: {momentum}") + if weight_decay < 0.0: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + + defaults = dict( + lr=lr, + momentum=momentum, + dampening=dampening, + weight_decay=weight_decay, + nesterov=nesterov, + trust_coefficient=trust_coefficient, + eps=eps, + ) + if nesterov and (momentum <= 0 or dampening != 0): + raise ValueError("Nesterov momentum requires a momentum and zero dampening") + + super().__init__(params, defaults) + + def __setstate__(self, state: Dict[str, Any]) -> None: + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("nesterov", False) + + # Type ignore for overloads is required for Python 3.7. + @overload # type: ignore[override] + def step(self, closure: None = None) -> None: ... + + @overload # type: ignore[override] + def step(self, closure: Callable[[], float]) -> float: ... + + @torch.no_grad() + def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]: + """Performs a single optimization step. + + Args: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + # Exclude scaling for params with 0 weight decay. + for group in self.param_groups: + weight_decay = group["weight_decay"] + momentum = group["momentum"] + dampening = group["dampening"] + nesterov = group["nesterov"] + + for p in group["params"]: + if p.grad is None: + continue + + d_p = p.grad + p_norm = torch.norm(p.data) + g_norm = torch.norm(p.grad.data) + + # Apply Lars scaling and weight decay. + if weight_decay != 0: + if p_norm != 0 and g_norm != 0: + lars_lr = p_norm / ( + g_norm + p_norm * weight_decay + group["eps"] + ) + lars_lr *= group["trust_coefficient"] + + d_p = d_p.add(p, alpha=weight_decay) + d_p *= lars_lr + + # Apply momentum. + if momentum != 0: + param_state = self.state[p] + if "momentum_buffer" not in param_state: + buf = param_state["momentum_buffer"] = torch.clone(d_p).detach() + else: + buf = param_state["momentum_buffer"] + buf.mul_(momentum).add_(d_p, alpha=1 - dampening) + + if nesterov: + d_p = d_p.add(buf, alpha=momentum) + else: + d_p = buf + + p.add_(d_p, alpha=-group["lr"]) + + return loss diff --git a/lightly/optim/param_groups.py b/lightly/optim/param_groups.py new file mode 100644 index 000000000..a09cba88b --- /dev/null +++ b/lightly/optim/param_groups.py @@ -0,0 +1,38 @@ +"""Optimiser parameter groups.""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from torch.nn import Module, Parameter + +from lightly.models.utils import get_weight_decay_parameters + +__all__ = ["param_groups"] + + +def param_groups(*modules: Module, weight_decay: float) -> List[Dict[str, Any]]: + """Splits parameters into a decayed group and a group that is not decayed. + + Normalization parameters and biases are the ones left out, which is what + every SSL reference implementation does and what the published numbers were + produced with. + + Example: + >>> optimizer = LARS(param_groups(backbone, head, weight_decay=1e-6), lr=4.8) + + Args: + *modules: The modules whose parameters to group. + weight_decay: The decay applied to everything except norms and biases. + + Returns: + Two groups, ready to hand to an optimiser. Both carry an explicit + ``weight_decay``, so the optimiser's own default never applies. + """ + decayed: List[Parameter] + not_decayed: List[Parameter] + decayed, not_decayed = get_weight_decay_parameters(modules) + return [ + {"name": "decay", "params": decayed, "weight_decay": weight_decay}, + {"name": "no_weight_decay", "params": not_decayed, "weight_decay": 0.0}, + ] diff --git a/lightly/optim/schedulers.py b/lightly/optim/schedulers.py new file mode 100644 index 000000000..77ae03a06 --- /dev/null +++ b/lightly/optim/schedulers.py @@ -0,0 +1,237 @@ +import warnings +from typing import Optional + +import numpy as np +import torch + + +def cosine_schedule( + step: int, + max_steps: int, + start_value: float, + end_value: float, + period: Optional[int] = None, +) -> float: + """Use cosine decay to gradually modify start_value to reach target end_value. + + Args: + step: + Current step number. + max_steps: + Total number of steps. + start_value: + Starting value. + end_value: + Target value. + period: + The number of steps over which the cosine function completes a full cycle. + If no period is provided, the scheduler will complete a half cycle over + max_steps. + + Returns: + Cosine decay value. + + """ + if step < 0: + raise ValueError(f"Current step number {step} can't be negative.") + if max_steps < 0: + raise ValueError(f"Total step number {max_steps} can't be negative.") + if period is None and step > max_steps: + warnings.warn( + f"Current step number {step} exceeds max_steps {max_steps}.", + category=RuntimeWarning, + ) + if period is not None and period <= 0: + raise ValueError(f"Period {period} must be >= 1") + + decay: float + if period is not None: # "cycle" based on period, if provided + decay = ( + end_value + - (end_value - start_value) * (np.cos(2 * np.pi * step / period) + 1) / 2 + ) + elif max_steps <= 1: + # Avoid division by zero + decay = end_value + elif step >= max_steps - 1: + # Special case for Pytorch Lightning which updates LR scheduler also for epoch + # after last training epoch. + decay = end_value + else: + decay = ( + end_value + - (end_value - start_value) + * (np.cos(np.pi * step / (max_steps - 1)) + 1) + / 2 + ) + # Cast to float as numpy operations result in np.float64. Checkpoints with + # np.float64 values cannot be loaded with torch.load(..., weights_only=True). + return float(decay) + + +def cosine_warmup_schedule( + step: int, + max_steps: int, + start_value: float, + end_value: float, + warmup_steps: int, + warmup_start_value: float, + warmup_end_value: Optional[float] = None, + period: Optional[int] = None, +) -> float: + """Use cosine decay to gradually modify start_value to reach target end_value. + + Uses linear warmup for the first warmup_steps steps. + + Args: + step: + Current step number. + max_steps: + Total number of steps. + start_value: + Starting value. + end_value: + Target value. + warmup_steps: + Number of steps for warmup. + warmup_start_value: + Starting value for warmup. + warmup_end_value: + Target value for warmup. Defaults to start_value. + period: + The number of steps over which the cosine function completes a full cycle. + If no period is provided, the scheduler will complete a half cycle over + max_steps - warmup_steps. + + Returns: + Cosine decay value. + """ + if warmup_steps < 0: + raise ValueError(f"Warmup steps {warmup_steps} can't be negative.") + if warmup_steps > max_steps: + raise ValueError(f"Warmup steps {warmup_steps} must be <= max_steps.") + if step > max_steps: + warnings.warn( + f"Current step number {step} exceeds max_steps {max_steps}.", + category=RuntimeWarning, + ) + + if warmup_end_value is None: + warmup_end_value = start_value + + if step < warmup_steps: + # Use step + 1 to reach warmup_end_value at end of warmup. This means that the + # initial warmup_start_value is skipped which is oftentimes desired when setting + # it to 0 as this would result in no parameter updates. + return ( + warmup_start_value + + (warmup_end_value - warmup_start_value) * (step + 1) / warmup_steps + ) + else: + max_steps = max_steps - (warmup_steps if period is None else 1) + return cosine_schedule( + step=step - warmup_steps, + max_steps=max_steps, + start_value=start_value, + end_value=end_value, + period=period, + ) + + +class CosineWarmupScheduler(torch.optim.lr_scheduler.LambdaLR): + """Cosine warmup scheduler for learning rate. + + Args: + optimizer: + Optimizer object to schedule the learning rate. + warmup_epochs: + Number of warmup epochs or steps. + max_epochs: + Total number of training epochs or steps. + last_epoch: + The index of last epoch or step. + start_value: + Starting learning rate. + end_value: + Target learning rate. + warmup_start_value: + Starting learning rate for warmup. + warmup_end_value: + Target learning rate for warmup. Defaults to start_value. + + Note: The `epoch` arguments do not necessarily have to be epochs. Any step or index + can be used. The naming follows the PyTorch convention to use `epoch` for the steps + in the scheduler. + """ + + def __init__( + self, + optimizer: torch.optim.Optimizer, + warmup_epochs: int, + max_epochs: int, + last_epoch: int = -1, + start_value: float = 1.0, + end_value: float = 0.001, + period: Optional[int] = None, + warmup_start_value: float = 0.0, + warmup_end_value: Optional[float] = None, + ) -> None: + self.warmup_epochs = warmup_epochs + self.max_epochs = max_epochs + self.start_value = start_value + self.end_value = end_value + self.period = period + self.warmup_start_value = warmup_start_value + self.warmup_end_value = warmup_end_value + + super().__init__( + optimizer=optimizer, + lr_lambda=self.scale_lr, + last_epoch=last_epoch, + ) + + def scale_lr(self, epoch: int) -> float: + """Scale learning rate according to the current epoch number. + + Args: + epoch: + Current epoch number. + + Returns: + Scaled learning rate. + + """ + return cosine_warmup_schedule( + step=epoch, + max_steps=self.max_epochs, + start_value=self.start_value, + end_value=self.end_value, + warmup_steps=self.warmup_epochs, + warmup_start_value=self.warmup_start_value, + warmup_end_value=self.warmup_end_value, + period=self.period, + ) + + +def linear_warmup_schedule( + step: int, + warmup_steps: int, + start_value: float, + end_value: float, +) -> float: + if warmup_steps < 0: + raise ValueError(f"Warmup steps {warmup_steps} can't be negative.") + if step < 0: + raise ValueError(f"Current step number {step} can't be negative.") + if start_value < 0: + raise ValueError(f"Start value {start_value} can't be negative.") + if end_value <= 0: + raise ValueError(f"End value {end_value} can't be non-positive.") + if start_value > end_value: + raise ValueError( + f"Start value {start_value} must be less than or equal to end value {end_value}." + ) + if step < warmup_steps: + return start_value + step / warmup_steps * (end_value - start_value) + else: + return end_value diff --git a/lightly/utils/lars.py b/lightly/utils/lars.py index 44bf43e7b..082460f73 100644 --- a/lightly/utils/lars.py +++ b/lightly/utils/lars.py @@ -1,166 +1,5 @@ -from typing import Any, Callable, Dict, Optional, overload +"""Moved to :mod:`lightly.optim.lars`. This path keeps working through 2.x.""" -import torch -from torch.optim.optimizer import Optimizer +from lightly.optim.lars import LARS - -class LARS(Optimizer): - """Extends SGD in PyTorch with LARS scaling from the paper "Large batch training of - Convolutional Networks" [0]. - - Implementation from PyTorch Lightning Bolts [1]. - - - [0]: https://arxiv.org/pdf/1708.03888.pdf - - [1]: https://github.com/Lightning-Universe/lightning-bolts/blob/2dfe45a4cf050f120d10981c45cfa2c785a1d5e6/pl_bolts/optimizers/lars.py#L1 - - Args: - params: - Iterable of parameters to optimize or dicts defining parameter groups. - lr: - Learning rate - momentum: - Momentum factor. - weight_decay: - Weight decay (L2 penalty). - dampening: - Dampening for momentum. - nesterov: - Enables Nesterov momentum. - trust_coefficient: - Trust coefficient for computing learning rate. - eps: - Eps for division denominator. - - Example: - >>> model = torch.nn.Linear(10, 1) - >>> input = torch.Tensor(10) - >>> target = torch.Tensor([1.]) - >>> loss_fn = lambda input, target: (input - target) ** 2 - >>> optimizer = LARS(model.parameters(), lr=0.1, momentum=0.9) - >>> optimizer.zero_grad() - >>> loss_fn(model(input), target).backward() - >>> optimizer.step() - - .. note:: - The application of momentum in the SGD part is modified according to - the PyTorch standards. LARS scaling fits into the equation in the - following fashion. - - .. math:: - \begin{aligned} - g_{t+1} & = \text{lars_lr} * (\beta * p_{t} + g_{t+1}), \\ - v_{t+1} & = \\mu * v_{t} + g_{t+1}, \\ - p_{t+1} & = p_{t} - \text{lr} * v_{t+1}, - \\end{aligned} - - where :math:`p`, :math:`g`, :math:`v`, :math:`\\mu` and :math:`\beta` denote the - parameters, gradient, velocity, momentum, and weight decay respectively. - The :math:`lars_lr` is defined by Eq. 6 in the paper. - The Nesterov version is analogously modified. - - .. warning:: - Parameters with weight decay set to 0 will automatically be excluded from - layer-wise LR scaling. This is to ensure consistency with papers like SimCLR - and BYOL. - """ - - def __init__( - self, - params: Any, - lr: float, - momentum: float = 0, - dampening: float = 0, - weight_decay: float = 0, - nesterov: bool = False, - trust_coefficient: float = 0.001, - eps: float = 1e-8, - ): - if lr < 0.0: - raise ValueError(f"Invalid learning rate: {lr}") - if momentum < 0.0: - raise ValueError(f"Invalid momentum value: {momentum}") - if weight_decay < 0.0: - raise ValueError(f"Invalid weight_decay value: {weight_decay}") - - defaults = dict( - lr=lr, - momentum=momentum, - dampening=dampening, - weight_decay=weight_decay, - nesterov=nesterov, - trust_coefficient=trust_coefficient, - eps=eps, - ) - if nesterov and (momentum <= 0 or dampening != 0): - raise ValueError("Nesterov momentum requires a momentum and zero dampening") - - super().__init__(params, defaults) - - def __setstate__(self, state: Dict[str, Any]) -> None: - super().__setstate__(state) - for group in self.param_groups: - group.setdefault("nesterov", False) - - # Type ignore for overloads is required for Python 3.7. - @overload # type: ignore[override] - def step(self, closure: None = None) -> None: ... - - @overload # type: ignore[override] - def step(self, closure: Callable[[], float]) -> float: ... - - @torch.no_grad() - def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]: - """Performs a single optimization step. - - Args: - closure (callable, optional): A closure that reevaluates the model - and returns the loss. - """ - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # Exclude scaling for params with 0 weight decay. - for group in self.param_groups: - weight_decay = group["weight_decay"] - momentum = group["momentum"] - dampening = group["dampening"] - nesterov = group["nesterov"] - - for p in group["params"]: - if p.grad is None: - continue - - d_p = p.grad - p_norm = torch.norm(p.data) - g_norm = torch.norm(p.grad.data) - - # Apply Lars scaling and weight decay. - if weight_decay != 0: - if p_norm != 0 and g_norm != 0: - lars_lr = p_norm / ( - g_norm + p_norm * weight_decay + group["eps"] - ) - lars_lr *= group["trust_coefficient"] - - d_p = d_p.add(p, alpha=weight_decay) - d_p *= lars_lr - - # Apply momentum. - if momentum != 0: - param_state = self.state[p] - if "momentum_buffer" not in param_state: - buf = param_state["momentum_buffer"] = torch.clone(d_p).detach() - else: - buf = param_state["momentum_buffer"] - buf.mul_(momentum).add_(d_p, alpha=1 - dampening) - - if nesterov: - d_p = d_p.add(buf, alpha=momentum) - else: - d_p = buf - - p.add_(d_p, alpha=-group["lr"]) - - return loss +__all__ = ["LARS"] diff --git a/lightly/utils/scheduler.py b/lightly/utils/scheduler.py index 77ae03a06..e2f7f9851 100644 --- a/lightly/utils/scheduler.py +++ b/lightly/utils/scheduler.py @@ -1,237 +1,15 @@ -import warnings -from typing import Optional - -import numpy as np -import torch - - -def cosine_schedule( - step: int, - max_steps: int, - start_value: float, - end_value: float, - period: Optional[int] = None, -) -> float: - """Use cosine decay to gradually modify start_value to reach target end_value. - - Args: - step: - Current step number. - max_steps: - Total number of steps. - start_value: - Starting value. - end_value: - Target value. - period: - The number of steps over which the cosine function completes a full cycle. - If no period is provided, the scheduler will complete a half cycle over - max_steps. - - Returns: - Cosine decay value. - - """ - if step < 0: - raise ValueError(f"Current step number {step} can't be negative.") - if max_steps < 0: - raise ValueError(f"Total step number {max_steps} can't be negative.") - if period is None and step > max_steps: - warnings.warn( - f"Current step number {step} exceeds max_steps {max_steps}.", - category=RuntimeWarning, - ) - if period is not None and period <= 0: - raise ValueError(f"Period {period} must be >= 1") - - decay: float - if period is not None: # "cycle" based on period, if provided - decay = ( - end_value - - (end_value - start_value) * (np.cos(2 * np.pi * step / period) + 1) / 2 - ) - elif max_steps <= 1: - # Avoid division by zero - decay = end_value - elif step >= max_steps - 1: - # Special case for Pytorch Lightning which updates LR scheduler also for epoch - # after last training epoch. - decay = end_value - else: - decay = ( - end_value - - (end_value - start_value) - * (np.cos(np.pi * step / (max_steps - 1)) + 1) - / 2 - ) - # Cast to float as numpy operations result in np.float64. Checkpoints with - # np.float64 values cannot be loaded with torch.load(..., weights_only=True). - return float(decay) - - -def cosine_warmup_schedule( - step: int, - max_steps: int, - start_value: float, - end_value: float, - warmup_steps: int, - warmup_start_value: float, - warmup_end_value: Optional[float] = None, - period: Optional[int] = None, -) -> float: - """Use cosine decay to gradually modify start_value to reach target end_value. - - Uses linear warmup for the first warmup_steps steps. - - Args: - step: - Current step number. - max_steps: - Total number of steps. - start_value: - Starting value. - end_value: - Target value. - warmup_steps: - Number of steps for warmup. - warmup_start_value: - Starting value for warmup. - warmup_end_value: - Target value for warmup. Defaults to start_value. - period: - The number of steps over which the cosine function completes a full cycle. - If no period is provided, the scheduler will complete a half cycle over - max_steps - warmup_steps. - - Returns: - Cosine decay value. - """ - if warmup_steps < 0: - raise ValueError(f"Warmup steps {warmup_steps} can't be negative.") - if warmup_steps > max_steps: - raise ValueError(f"Warmup steps {warmup_steps} must be <= max_steps.") - if step > max_steps: - warnings.warn( - f"Current step number {step} exceeds max_steps {max_steps}.", - category=RuntimeWarning, - ) - - if warmup_end_value is None: - warmup_end_value = start_value - - if step < warmup_steps: - # Use step + 1 to reach warmup_end_value at end of warmup. This means that the - # initial warmup_start_value is skipped which is oftentimes desired when setting - # it to 0 as this would result in no parameter updates. - return ( - warmup_start_value - + (warmup_end_value - warmup_start_value) * (step + 1) / warmup_steps - ) - else: - max_steps = max_steps - (warmup_steps if period is None else 1) - return cosine_schedule( - step=step - warmup_steps, - max_steps=max_steps, - start_value=start_value, - end_value=end_value, - period=period, - ) - - -class CosineWarmupScheduler(torch.optim.lr_scheduler.LambdaLR): - """Cosine warmup scheduler for learning rate. - - Args: - optimizer: - Optimizer object to schedule the learning rate. - warmup_epochs: - Number of warmup epochs or steps. - max_epochs: - Total number of training epochs or steps. - last_epoch: - The index of last epoch or step. - start_value: - Starting learning rate. - end_value: - Target learning rate. - warmup_start_value: - Starting learning rate for warmup. - warmup_end_value: - Target learning rate for warmup. Defaults to start_value. - - Note: The `epoch` arguments do not necessarily have to be epochs. Any step or index - can be used. The naming follows the PyTorch convention to use `epoch` for the steps - in the scheduler. - """ - - def __init__( - self, - optimizer: torch.optim.Optimizer, - warmup_epochs: int, - max_epochs: int, - last_epoch: int = -1, - start_value: float = 1.0, - end_value: float = 0.001, - period: Optional[int] = None, - warmup_start_value: float = 0.0, - warmup_end_value: Optional[float] = None, - ) -> None: - self.warmup_epochs = warmup_epochs - self.max_epochs = max_epochs - self.start_value = start_value - self.end_value = end_value - self.period = period - self.warmup_start_value = warmup_start_value - self.warmup_end_value = warmup_end_value - - super().__init__( - optimizer=optimizer, - lr_lambda=self.scale_lr, - last_epoch=last_epoch, - ) - - def scale_lr(self, epoch: int) -> float: - """Scale learning rate according to the current epoch number. - - Args: - epoch: - Current epoch number. - - Returns: - Scaled learning rate. - - """ - return cosine_warmup_schedule( - step=epoch, - max_steps=self.max_epochs, - start_value=self.start_value, - end_value=self.end_value, - warmup_steps=self.warmup_epochs, - warmup_start_value=self.warmup_start_value, - warmup_end_value=self.warmup_end_value, - period=self.period, - ) - - -def linear_warmup_schedule( - step: int, - warmup_steps: int, - start_value: float, - end_value: float, -) -> float: - if warmup_steps < 0: - raise ValueError(f"Warmup steps {warmup_steps} can't be negative.") - if step < 0: - raise ValueError(f"Current step number {step} can't be negative.") - if start_value < 0: - raise ValueError(f"Start value {start_value} can't be negative.") - if end_value <= 0: - raise ValueError(f"End value {end_value} can't be non-positive.") - if start_value > end_value: - raise ValueError( - f"Start value {start_value} must be less than or equal to end value {end_value}." - ) - if step < warmup_steps: - return start_value + step / warmup_steps * (end_value - start_value) - else: - return end_value +"""Moved to :mod:`lightly.optim.schedulers`. This path keeps working through 2.x.""" + +from lightly.optim.schedulers import ( + CosineWarmupScheduler, + cosine_schedule, + cosine_warmup_schedule, + linear_warmup_schedule, +) + +__all__ = [ + "CosineWarmupScheduler", + "cosine_schedule", + "cosine_warmup_schedule", + "linear_warmup_schedule", +] diff --git a/tests/optim/__init__.py b/tests/optim/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/optim/test_param_groups.py b/tests/optim/test_param_groups.py new file mode 100644 index 000000000..14779392c --- /dev/null +++ b/tests/optim/test_param_groups.py @@ -0,0 +1,62 @@ +from typing import Set + +import torch +from torch.nn import BatchNorm1d, Conv2d, LayerNorm, Linear, Module, Sequential + +from lightly.optim import LARS, param_groups + + +class Model(Module): + def __init__(self) -> None: + super().__init__() + self.conv = Conv2d(3, 4, kernel_size=3) + self.norm = BatchNorm1d(4) + self.head = Sequential(Linear(4, 4), LayerNorm(4)) + + +def names_of(model: Module, wanted: Set[int]) -> Set[str]: + return {name for name, p in model.named_parameters() if id(p) in wanted} + + +def test_norms_and_biases_are_not_decayed() -> None: + model = Model() + decay, no_decay = param_groups(model, weight_decay=1e-6) + + assert decay["weight_decay"] == 1e-6 + assert no_decay["weight_decay"] == 0.0 + assert names_of(model, {id(p) for p in decay["params"]}) == { + "conv.weight", + "head.0.weight", + } + assert names_of(model, {id(p) for p in no_decay["params"]}) == { + "conv.bias", + "norm.weight", + "norm.bias", + "head.0.bias", + "head.1.weight", + "head.1.bias", + } + + +def test_every_parameter_lands_in_exactly_one_group() -> None: + model = Model() + groups = param_groups(model, weight_decay=1e-6) + grouped = [id(p) for group in groups for p in group["params"]] + assert sorted(grouped) == sorted(id(p) for p in model.parameters()) + + +def test_several_modules_are_grouped_together() -> None: + backbone, head = Model(), Linear(4, 2) + decay, no_decay = param_groups(backbone, head, weight_decay=0.1) + assert any(p is head.weight for p in decay["params"]) + assert any(p is head.bias for p in no_decay["params"]) + + +def test_the_groups_are_what_an_optimiser_takes() -> None: + model = Model() + optimizer = LARS(param_groups(model, weight_decay=1e-6), lr=0.1, momentum=0.9) + assert [group["weight_decay"] for group in optimizer.param_groups] == [1e-6, 0.0] + linear = model.head[0] + assert isinstance(linear, Linear) + linear.weight.grad = torch.ones_like(linear.weight) + optimizer.step()