torchzero 0.1.8__py3-none-any.whl → 0.3.2__py3-none-any.whl

torchzero/modules/scheduling/step_size.py

@@ -1,80 +0,0 @@
-import random
-from typing import Any
-
-from ...core import OptimizerModule
-from ...tensorlist import TensorList
-
-
-class PolyakStepSize(OptimizerModule):
-    """Polyak step-size. Meant to be used at the beginning when ascent is the gradient but other placements may work.
-    This can also work with SGD as SPS (Stochastic Polyak Step-Size) seems to use the same formula.
-
-    Args:
-        max (float | None, optional): maximum possible step size. Defaults to None.
-        min_obj_value (int, optional): (estimated) minimal possible value of the objective function (lowest possible loss). Defaults to 0.
-        use_grad (bool, optional):
-            if True, uses dot product of update and gradient to compute the step size.
-            Otherwise, dot product of update with itself is used, which has no geometric meaning so it probably won't work well.
-            Defaults to True.
-        parameterwise (bool, optional):
-            if True, calculate Polyak step-size for each parameter separately,
-            if False calculate one global step size for all parameters. Defaults to False.
-        alpha (float, optional): multiplier to Polyak step-size. Defaults to 1.
-    """
-    def __init__(self, max: float | None = None, min_obj_value: float = 0, use_grad=True, parameterwise=False, alpha: float = 1):
-
-        defaults = dict(alpha = alpha)
-        super().__init__(defaults)
-        self.max = max
-        self.min_obj_value = min_obj_value
-        self.use_grad = use_grad
-        self.parameterwise = parameterwise
-
-    def _update(self, vars, ascent):
-        if vars.closure is None: raise ValueError("PolyakStepSize requires closure")
-        if vars.fx0 is None: vars.fx0 = vars.closure(False) # can only happen when placed after SPSA
-
-        alpha = self.get_group_key('alpha')
-
-        if self.parameterwise:
-            if self.use_grad: denom = (ascent*vars.maybe_compute_grad_(self.get_params())).mean()
-            else: denom = ascent.pow(2).mean()
-            polyak_step_size: TensorList | Any = (vars.fx0 - self.min_obj_value) / denom.where(denom!=0, 1) # type:ignore
-            polyak_step_size = polyak_step_size.where(denom != 0, 0)
-            if self.max is not None: polyak_step_size = polyak_step_size.clamp_max(self.max)
-
-        else:
-            if self.use_grad: denom = (ascent*vars.maybe_compute_grad_(self.get_params())).total_mean()
-            else: denom = ascent.pow(2).total_mean()
-            if denom == 0: polyak_step_size = 0 # we converged
-            else: polyak_step_size = (vars.fx0 - self.min_obj_value) / denom
-
-            if self.max is not None:
-                if polyak_step_size > self.max: polyak_step_size = self.max
-
-        ascent.mul_(alpha * polyak_step_size)
-        return ascent
-
-
-
-class RandomStepSize(OptimizerModule):
-    """Uses random global step size from `low` to `high`.
-
-    Args:
-        low (float, optional): minimum learning rate. Defaults to 0.
-        high (float, optional): maximum learning rate. Defaults to 1.
-        parameterwise (bool, optional):
-            if True, generate random step size for each parameter separately,
-            if False generate one global random step size. Defaults to False.
-    """
-    def __init__(self, low: float = 0, high: float = 1, parameterwise=False):
-        super().__init__({})
-        self.low = low; self.high = high
-        self.parameterwise = parameterwise
-
-    def _update(self, vars, ascent):
-        if self.parameterwise:
-            lr = [random.uniform(self.low, self.high) for _ in range(len(ascent))]
-        else:
-            lr = random.uniform(self.low, self.high)
-        return ascent.mul_(lr) # type:ignore

torchzero/modules/smoothing/gaussian_smoothing.py

@@ -1,90 +0,0 @@
-from contextlib import nullcontext
-
-import numpy as np
-import torch
-
-from ...tensorlist import TensorList, Distributions, mean as tlmean
-from ...utils.python_tools import _ScalarLoss
-from ...core import _ClosureType, OptimizationVars, OptimizerModule, _maybe_pass_backward
-
-
-def _numpy_or_torch_mean(losses: list):
-    """Returns the mean of a list of losses, which can be either numpy arrays or torch tensors."""
-    if isinstance(losses[0], torch.Tensor):
-        return torch.mean(torch.stack(losses))
-    return np.mean(losses).item()
-
-class GaussianHomotopy(OptimizerModule):
-    """Samples and averages value and gradients in multiple random points around current position.
-    This effectively applies smoothing to the function.
-
-    Args:
-        n_samples (int, optional): number of gradient samples from around current position. Defaults to 4.
-        sigma (float, optional): how far from current position to sample from. Defaults to 0.1.
-        distribution (tl.Distributions, optional): distribution for random positions. Defaults to "normal".
-        sample_x0 (bool, optional): 1st sample will be x0. Defaults to False.
-        randomize_every (int | None, optional): randomizes the points every n steps. Defaults to 1.
-    """
-    def __init__(
-        self,
-        n_samples: int = 4,
-        sigma: float = 0.1,
-        distribution: Distributions = "normal",
-        sample_x0 = False,
-        randomize_every: int | None = 1,
-    ):
-        defaults = dict(sigma = sigma)
-        super().__init__(defaults)
-        self.n_samples = n_samples
-        self.distribution: Distributions = distribution
-        self.randomize_every = randomize_every
-        self.current_step = 0
-        self.perturbations = None
-        self.sample_x0 = sample_x0
-
-    @torch.no_grad()
-    def step(self, vars: OptimizationVars):
-        if vars.closure is None: raise ValueError('GaussianSmoothing requires closure.')
-        closure = vars.closure
-        params = self.get_params()
-        sigmas = self.get_group_key('sigma')
-
-        # generate random perturbations
-        if self.perturbations is None or (self.randomize_every is not None and self.current_step % self.randomize_every == 0):
-            if self.sample_x0:
-                self.perturbations = [params.sample_like(sigmas, distribution=self.distribution) for _ in range(self.n_samples-1)]
-            else:
-                self.perturbations = [params.sample_like(sigmas, distribution=self.distribution) for _ in range(self.n_samples)]
-
-        @torch.no_grad
-        def smooth_closure(backward = True):
-            losses = []
-            grads = []
-
-            # sample gradient and loss at x0
-            if self.sample_x0:
-                with torch.enable_grad() if backward else nullcontext():
-                    losses.append(closure())
-                    if backward: grads.append(params.grad.clone())
-
-            # sample gradients from points around current params
-            # and average them
-            if self.perturbations is None: raise ValueError('who set perturbations to None???')
-            for p in self.perturbations:
-                params.add_(p)
-                with torch.enable_grad() if backward else nullcontext():
-                    losses.append(_maybe_pass_backward(closure, backward))
-                    if backward: grads.append(params.grad.clone())
-                params.sub_(p)
-
-            # set the new averaged grads and return average loss
-            if backward: params.set_grad_(tlmean(grads))
-            return _numpy_or_torch_mean(losses)
-
-
-        self.current_step += 1
-        vars.closure = smooth_closure
-        return self._update_params_or_step_with_next(vars)
-
-
-# todo single loop gaussian homotopy?

torchzero/modules/weight_averaging/__init__.py

@@ -1,2 +0,0 @@
-from .ema import SwitchEMA
-from .swa import PeriodicSWA, CyclicSWA

torchzero/modules/weight_averaging/ema.py

@@ -1,72 +0,0 @@
-import torch
-from ...core import OptimizerModule
-
-
-def _reset_stats_hook(optimizer, state):
-    for module in optimizer.unrolled_modules:
-        module: OptimizerModule
-        module.reset_stats()
-
-# the reason why this needs to be at the end is ??? I NEED TO REMEMBER
-class SwitchEMA(OptimizerModule):
-    """Switch-EMA. Every n steps switches params to an exponential moving average of past weights.
-
-    In the paper the switch happens after each epoch.
-
-    Please put this module at the end, after all other modules.
-
-    This can also function as EMA, set `update_every` to None and instead call `set_ema` and `unset_ema` on this module.
-
-
-    Args:
-        update_every (int): number of steps (batches) between setting model parameters to EMA.
-        momentum (int): EMA momentum factor.
-        reset_stats (bool, optional):
-            if True, when setting model parameters to EMA, resets other modules stats such as momentum velocities.
-            It might be better to set this to False if `update_every` is very small. Defaults to True.
-
-    reference
-        https://arxiv.org/abs/2402.09240
-    """
-    def __init__(self, update_every: int | None, momentum: float = 0.99, reset_stats: bool = True):
-        defaults = dict(momentum=momentum)
-        super().__init__(defaults)
-        self.update_every = update_every
-        self.cur_step = 0
-        self.update_every = update_every
-        self._reset_stats = reset_stats
-        self.orig_params = None
-
-    def set_ema(self):
-        """sets module parameters to EMA, stores original parameters that can be restored by calling `unset_ema`"""
-        params = self.get_params()
-        self.orig_params = params.clone()
-        params.set_(self.get_state_key('ema', init = 'params', params=params))
-
-    def unset_ema(self):
-        """Undoes `set_ema`."""
-        if self.orig_params is None: raise ValueError('call `set_ema` first, and then `unset_ema`.')
-        params = self.get_params()
-        params.set_(self.orig_params)
-
-    @torch.no_grad
-    def step(self, vars):
-        # if self.next_module is not None:
-        #     warn(f'EMA should usually be the last module, but {self.next_module.__class__.__name__} is after it.')
-        self.cur_step += 1
-
-        params = self.get_params()
-        # state.maybe_use_grad_(params)
-        # update params with the child. Averaging is always applied at the end.
-        ret = self._update_params_or_step_with_next(vars, params)
-
-        ema = self.get_state_key('ema', init = 'params', params=params)
-        momentum = self.get_group_key('momentum')
-
-        ema.lerp_compat_(params, 1 - momentum)
-
-        if (self.update_every is not None) and (self.cur_step % self.update_every == 0):
-            params.set_(ema.clone())
-            if self._reset_stats: vars.add_post_step_hook(_reset_stats_hook)
-
-        return ret

torchzero/modules/weight_averaging/swa.py

@@ -1,171 +0,0 @@
-from ...core import OptimizerModule
-
-
-def _reset_stats_hook(optimizer, state):
-    for module in optimizer.unrolled_modules:
-        module: OptimizerModule
-        module.reset_stats()
-
-class PeriodicSWA(OptimizerModule):
-    """Periodic Stochastic Weight Averaging.
-
-    Please put this module at the end, after all other modules.
-
-    The algorithm is as follows:
-
-    1. perform `pswa_start` normal steps before starting PSWA.
-
-    2. Perform multiple SWA iterations. On each iteration,
-    run SWA algorithm for `num_cycles` cycles,
-    and set weights to the weighted average before starting the next SWA iteration.
-
-    SWA iteration is as follows:
-
-    1. perform `cycle_start` initial steps (can be 0)
-
-    2. for `num_cycles`, after every `cycle_length` steps passed, update the weight average with current model weights.
-
-    3. After `num_cycles` cycles passed, set model parameters to the weight average.
-
-    Args:
-        first_swa (int):
-            number of steps before starting PSWA, authors run PSWA starting from 40th epoch out ot 150 epochs in total.
-        cycle_length (int):
-            number of steps betwen updating the weight average. Authors update it once per epoch.
-        num_cycles (int):
-            Number of weight average updates before setting model weights to the average and proceding to the next cycle.
-            Authors use 20 (meaning 20 epochs since each cycle is 1 epoch).
-        cycle_start (int, optional):
-            number of steps at the beginning of each SWA period before updating the weight average (default: 0).
-        reset_stats (bool, optional):
-            if True, when setting model parameters to SWA, resets other modules stats such as momentum velocities (default: True).
-    """
-    def __init__(self, pswa_start: int, cycle_length: int, num_cycles: int, cycle_start: int = 0, reset_stats:bool = True):
-
-        super().__init__({})
-        self.pswa_start = pswa_start
-        self.cycle_start = cycle_start
-        self.cycle_length = cycle_length
-        self.num_cycles = num_cycles
-        self._reset_stats = reset_stats
-
-
-        self.cur = 0
-        self.period_cur = 0
-        self.swa_cur = 0
-        self.n_models = 0
-
-    def step(self, vars):
-        swa = None
-        params = self.get_params()
-        ret = self._update_params_or_step_with_next(vars, params)
-
-        # start first period after `pswa_start` steps
-        if self.cur >= self.pswa_start:
-
-            # start swa after `cycle_start` steps in the current period
-            if self.period_cur >= self.cycle_start:
-
-                # swa updates on every `cycle_length`th step
-                if self.swa_cur % self.cycle_length == 0:
-                    swa = self.get_state_key('swa') # initialized to zeros for simplicity
-                    swa.mul_(self.n_models).add_(params).div_(self.n_models + 1)
-                    self.n_models += 1
-
-                self.swa_cur += 1
-
-            self.period_cur += 1
-
-        self.cur += 1
-
-        # passed num_cycles in period, set model parameters to SWA
-        if self.n_models == self.num_cycles:
-            self.period_cur = 0
-            self.swa_cur = 0
-            self.n_models = 0
-
-            assert swa is not None # it's created above self.n_models += 1
-
-            params.set_(swa)
-            # add a hook that resets momentum, which also deletes `swa` in this module
-            if self._reset_stats: vars.add_post_step_hook(_reset_stats_hook)
-
-        return ret
-
-class CyclicSWA(OptimizerModule):
-    """Periodic SWA with cyclic learning rate. So it samples the weights, increases lr to `peak_lr`, samples the weights again,
-    decreases lr back to `init_lr`, and samples the weights last time. Then model weights are replaced with the average of the three sampled weights,
-    and next cycle starts. I made this due to a horrible misreading of the original SWA paper but it seems to work well.
-
-    Please put this module at the end, after all other modules.
-
-    Args:
-        cswa_start (int): number of steps before starting the first CSWA cycle.
-        cycle_length (int): length of each cycle in steps.
-        steps_between (int): number of steps between cycles.
-        init_lr (float, optional): initial and final learning rate in each cycle. Defaults to 0.
-        peak_lr (float, optional): peak learning rate of each cycle. Defaults to 1.
-        sample_all (float, optional): if True, instead of sampling 3 weights, it samples all weights in the cycle. Defaults to False.
-        reset_stats (bool, optional):
-            if True, when setting model parameters to SWA, resets other modules stats such as momentum velocities (default: True).
-
-    """
-    def __init__(self, cswa_start: int, cycle_length: int, steps_between: int, init_lr: float = 0, peak_lr: float = 1, sample_all = False, reset_stats: bool=True,):
-        defaults = dict(init_lr = init_lr, peak_lr = peak_lr)
-        super().__init__(defaults)
-        self.cswa_start = cswa_start
-        self.cycle_length = cycle_length
-        self.init_lr = init_lr
-        self.peak_lr = peak_lr
-        self.steps_between = steps_between
-        self.sample_all = sample_all
-        self._reset_stats = reset_stats
-
-        self.cur = 0
-        self.cycle_cur = 0
-        self.n_models = 0
-
-        self.cur_lr = self.init_lr
-
-    def step(self, vars):
-        params = self.get_params()
-
-        # start first period after `cswa_start` steps
-        if self.cur >= self.cswa_start:
-
-            ascent = vars.maybe_use_grad_(params)
-
-            # determine the lr
-            point = self.cycle_cur / self.cycle_length
-            init_lr, peak_lr = self.get_group_keys('init_lr', 'peak_lr')
-            if point < 0.5:
-                p2 = point*2
-                lr = init_lr * (1-p2) + peak_lr * p2
-            else:
-                p2 = (1 - point)*2
-                lr = init_lr * (1-p2) + peak_lr * p2
-
-            ascent *= lr
-            ret = self._update_params_or_step_with_next(vars, params)
-
-            if self.sample_all or self.cycle_cur in (0, self.cycle_length, self.cycle_length // 2):
-                swa = self.get_state_key('swa')
-                swa.mul_(self.n_models).add_(params).div_(self.n_models + 1)
-                self.n_models += 1
-
-                if self.cycle_cur == self.cycle_length:
-                    if not self.sample_all: assert self.n_models == 3, self.n_models
-                    self.n_models = 0
-                    self.cycle_cur = -1
-
-                    params.set_(swa)
-                    if self._reset_stats: vars.add_post_step_hook(_reset_stats_hook)
-
-            self.cycle_cur += 1
-
-        else:
-            ret = self._update_params_or_step_with_next(vars, params)
-
-        self.cur += 1
-
-        return ret

torchzero/optim/experimental/__init__.py

@@ -1,20 +0,0 @@
-"""Optimizers that I haven't tested and various (mostly stupid) ideas go there.
-If something works well I will move it outside of experimental folder.
-Otherwise all optimizers in this category should be considered unlikely to good for most tasks."""
-from .experimental import (
-    HVPDiagNewton,
-    ExaggeratedNesterov,
-    ExtraCautiousAdam,
-    GradMin,
-    InwardSGD,
-    MinibatchRprop,
-    MomentumDenominator,
-    MomentumNumerator,
-    MultistepSGD,
-    RandomCoordinateMomentum,
-    ReciprocalSGD,
-    NoiseSign,
-)
-
-
-from .ray_search import NewtonFDMRaySearch, LBFGSRaySearch