gradboard 5.1.0__py3-none-any.whl

gradboard/__init__.py

@@ -0,0 +1,3 @@
+from . import cycles
+from . import optimiser
+from . import scheduler

gradboard/cycles.py

@@ -0,0 +1,319 @@
+"""
+Utilities for generating a range of learning rate schedules.
+"""
+
+import math
+from typing import Optional, List, Union, Callable
+
+
+def ascent(step: int, total_steps: int) -> float:
+    """
+    Get a sequence of numbers evenly spaced between 0 and 1 so that the first
+        number is 0 and the last is 1 and there are `total_steps` numbers in
+        the sequence.
+    """
+    return round(step / (total_steps - 0.999), 8)
+
+
+def triangle(step: int, total_steps: int) -> float:
+    """
+    Get a triangular sequence of numbers between 0 and 1, going up in half of
+        `total_steps` and coming down in the other half, peaking at ~1.
+    """
+    half = int(math.ceil(total_steps / 2))
+    if step < half:
+        return 2 * ascent(step, total_steps)
+    else:
+        return 2 - 2 * ascent(step, total_steps)
+
+
+def cosine(step: int, total_steps: int) -> float:
+    """
+    Get a sequence of numbers between 0 and 1 in the shape of a cosine wave with
+        wavelength `total_steps`.
+    """
+    if total_steps < 1:
+        raise ValueError(f"total_steps must be >= 1, got {total_steps}")
+    angle = (step / (total_steps - 0.999)) * (2 * math.pi)
+    return round((math.cos(angle) + 1) / 2, 8)
+
+
+def quarter_circle(step: int, total_steps: int) -> float:
+    """
+    Get a sequence of numbers between 0 and 1 in the shape of a quarter-circle with
+        radius `total_steps'.
+    """
+    if total_steps < 1:
+        raise ValueError(f"total_steps must be >= 1, got {total_steps}")
+    x = 0 if total_steps == 1 else step / (total_steps - 1)
+    # x^2 + y^2 = r^2 = 1
+    # Therefore y^2 = 1 - x^2
+    # Therefore y^2 = (1 + x)(1 - x)
+    y_squared = max(1 - x, 0) * (1 + x)
+    return math.sqrt(y_squared)
+
+
+def half_cosine(step: int, total_steps: int) -> float:
+    """
+    Get a sequence of numbers between 0 and 1 in the shape of the descending
+        half of a cosine wave with wavelength 2*`total_steps`.
+    """
+    return cosine(step, (total_steps * 2) - 1)
+
+
+def cycloid(step: int, total_steps: int) -> float:
+    """
+    Get a sequence of numbers between 0 and 1 in the shape of a cycloid with
+        circle diameter 1.0 and `total_steps/(2*math.pi)` steps per cycle.
+    """
+    x = step * (math.pi / (total_steps - 1))
+
+    def fx(t):
+        return 0.5 * (t - math.sin(t)) - x
+
+    def fx_prime(t):
+        return 0.5 - 0.5 * math.cos(t)
+
+    def fy_prime(t):
+        return 0.5 - 0.5 * -math.sin(t)
+
+    angle_estimate = 0.5 * x
+
+    # XXX: 200 iterations is too many! Use a more efficient root finding algorithm
+    for _ in range(200):
+        if abs(fx_prime(angle_estimate)) > 0.1:
+            update = fx(angle_estimate) / fx_prime(angle_estimate)
+        else:
+            update = fx(angle_estimate) / fy_prime(angle_estimate)
+        angle_estimate = angle_estimate - update
+
+    return 0.5 * (1 - math.cos(angle_estimate))
+
+
+def half_cycloid(step: int, total_steps: int) -> float:
+    return cycloid(total_steps + step, 2 * total_steps)
+
+
+FN_LIBRARY = {
+    "ascent": ascent,
+    "triangle": triangle,
+    "cosine": cosine,
+    "half_cosine": half_cosine,
+    "quarter_circle": quarter_circle,
+    "half_cycloid": half_cycloid,
+}
+
+
+class Cycle:
+    def __init__(
+        self,
+        generating_function: Union[str, Callable],
+        training_examples,
+        epochs,
+        batch_size,
+        t_0: Optional[int] = None,
+        t_mult: float = 1.0,
+        t_scale: float = 1.0,
+        low=0.0,
+        high=1.0,
+        reflect=False,
+    ):
+        self.training_examples = training_examples
+        self.epochs = epochs
+        self.batch_size = batch_size
+        self.total_steps = int(
+            epochs * (math.floor(training_examples / batch_size) + 1)
+        )
+
+        self.t_0 = (
+            t_0 * (training_examples / batch_size)
+            if t_0 is not None
+            else self.total_steps
+        )
+        self.t_mult = t_mult
+        self.t_scale = t_scale
+
+        self.low = low
+        self.high = high
+
+        self.reflect = reflect
+
+        if not callable(generating_function):
+            if generating_function in FN_LIBRARY:
+                self._generating_function = FN_LIBRARY[generating_function]
+            else:
+                raise NotImplementedError(
+                    "`generating_function` must be a callable object or one of "
+                    '"ascent", "triangle", "cosine", "half_cosine", "quarter_circle" '
+                    'or "half_cycloid"'
+                )
+        else:
+            self._generating_function = generating_function
+
+    def _get_window(self, step):
+        windows = self._windows()
+        cumulative = [
+            sum([w[0] for w in windows][: i + 1]) for i in range(len(windows))
+        ]
+        position = None
+        local_step = None
+        for i, c in enumerate(cumulative):
+            if c > step:
+                position = i
+                local_step = step if i == 0 else step - cumulative[i - 1]
+                break
+        window_width, window_height = windows[position]
+        return window_width, local_step, window_height
+
+    def _generate(self, step) -> list:
+        total_steps, step, scale = self._get_window(step)
+        y = self._generating_function(step, total_steps)
+        y = y * scale
+        y = 1 - y if self.reflect else y
+        return y * (self.high - self.low) + self.low
+
+    def __call__(self, n):
+        return self._generate(n)
+
+    def __len__(self):
+        return self.total_steps
+
+    def _windows(self):
+        assert self.t_mult > 0
+
+        # Get tile widths
+        widths = [self.t_0]
+        while True:
+            next_item = widths[-1] * self.t_mult
+            if sum(widths) + next_item <= self.total_steps:
+                widths.append(next_item)
+            else:
+                break
+        for i in range(1, len(widths)):
+            widths[i] = int(widths[i] * (self.total_steps / sum(widths)))
+        widths[-1] += self.total_steps - sum(widths)
+
+        # Get tile heights
+        heights = [1.0 * self.t_scale**i for i in range(len(widths))]
+
+        return list(zip(widths, heights, strict=True))
+
+    @property
+    def stats(self) -> float:
+        """
+        Returns the area (as a percentage of the area of a curve where the learning
+            rate is constant max_lr), percentage ascent steps and percentage descent
+            steps of a learning rate schedule.
+        """
+        total_area = 0
+        max_area = 0
+        ascent_steps = 0
+        descent_steps = 0
+        avg_up_gradient = 0
+        avg_down_gradient = 0
+        total_gradient = 0
+        previous_lr = None
+        for s in range(self.total_steps):
+            height = self(s)
+            total_area += height
+            max_area += 1
+            if previous_lr is None:
+                pass
+            elif previous_lr > height:
+                descent_steps += 1
+                avg_down_gradient += height - previous_lr
+                total_gradient += height - previous_lr
+            elif previous_lr < height:
+                ascent_steps += 1
+                avg_up_gradient += height - previous_lr
+                total_gradient += height - previous_lr
+            else:
+                total_gradient += height
+            previous_lr = height
+        return {
+            "area": total_area / max_area,
+            "pc_ascent": round(ascent_steps / self.total_steps, 3),
+            "pc_descent": round(descent_steps / self.total_steps, 3),
+            "avg_up_gradient": round(avg_up_gradient, 3),
+            "avg_down_gradient": round(avg_down_gradient, 3),
+            "avg_gradient": round(-(self.high - self.low) / self.total_steps, 3),
+        }
+
+
+class CycleProduct(Cycle):
+    def __init__(self, cycles: List[Cycle], reflect=False, normalise: bool = False):
+        """
+        Args:
+            normalise: if true, the square root of the product is returned (i.e.
+                the geometric mean of the two cycles that were multiplied together)
+        """
+        main_training_examples = cycles[0].training_examples
+        main_batch_size = cycles[0].batch_size
+
+        assert all(c.training_examples == main_training_examples for c in cycles)
+        assert all(c.batch_size == main_batch_size for c in cycles)
+
+        self.cycles = cycles
+        self.reflect = reflect
+        self.normalise = normalise
+
+        def generating_function(step: int, total_steps: int) -> float:
+            output = self.cycles[0](step)
+            for c in self.cycles[1:]:
+                output *= c(step % c.total_steps)
+            if self.normalise:
+                output = math.sqrt(output)
+            return output
+
+        super().__init__(
+            generating_function=generating_function,
+            training_examples=self.cycles[0].training_examples,
+            epochs=self.cycles[0].epochs,
+            batch_size=self.cycles[0].batch_size,
+            reflect=reflect,
+        )
+
+
+class CycleSequence:
+    def __init__(self, cycles: List[Cycle]):
+        self.total_steps = sum([c.total_steps for c in cycles])
+        self.cycles = cycles
+
+    def _generate(self, step):
+        cycle, step = self._get_cycle_and_step(step)
+        return self.cycles[cycle](step)
+
+    def _get_cycle_and_step(self, step):
+        cycle_lengths = [c.total_steps for c in self.cycles]
+        cumulative = [sum(cycle_lengths[: i + 1]) for i in range(len(cycle_lengths))]
+        cycle = None
+        local_step = None
+        for i, c in enumerate(cumulative):
+            if c > step:
+                cycle = i
+                local_step = step if i == 0 else step - cumulative[i - 1]
+                break
+        return cycle, local_step
+
+    def __call__(self, step):
+        return self._generate(step)
+
+    def __len__(self):
+        return self.total_steps
+
+    @property
+    def stats(self) -> float:
+        """
+        Returns the area (as a percentage of the area of a curve where the learning
+            rate is constant max_lr), percentage ascent steps and percentage descent
+            steps of a learning rate schedule.
+        """
+        cycle_ratios = [c.total_steps for c in self.cycles]
+        cycle_stats = {k: v * cycle_ratios[0] for k, v in self.cycles[0].stats.items()}
+        for i, cycle in enumerate(self.cycles):
+            if i == 0:
+                continue  # We already did the first one, above
+            for k, v in cycle.stats.items():
+                cycle_stats[k] += cycle_ratios[i] * v
+
+        return {k: v / self.total_steps for k, v in cycle_stats.items()}

gradboard/optimiser.py

@@ -0,0 +1,48 @@
+import math
+import warnings
+import torch
+from torch.optim.optimizer import Optimizer
+from torch.optim import AdamW
+
+EXCLUDE_FROM_WEIGHT_DECAY = ["nondecay", "bias", "norm", "embedding", "beta"]
+
+
+def get_optimiser(
+    model,
+    optimiser=AdamW,
+    lr=1e-3,
+    weight_decay=1e-2,
+    exclude_keywords=EXCLUDE_FROM_WEIGHT_DECAY,
+):
+    """
+    Defaults are from one of the presets from the accompanying repo to Hassani
+        et al. (2023) "Escaping the Big Data Paradigm with Compact Transformers",
+        https://github.com/SHI-Labs/Compact-Transformers/blob/main/configs/
+        pretrained/cct_7-3x1_cifar100_1500epochs.yml
+    """
+    weight_decay_exclude = []
+
+    for keyword in exclude_keywords:
+        weight_decay_exclude += [
+            p for name, p in model.named_parameters() if keyword in name.lower()
+        ]
+
+    weight_decay_exclude = set(weight_decay_exclude)
+
+    if len(weight_decay_exclude) > 0:
+        warnings.warn(
+            "Excluded the following parameters from weight decay based on "
+            "exclude keywords: {weight_decay_exclude}",
+            stacklevel=2,
+        )
+
+    weight_decay_include = set(model.parameters()) - weight_decay_exclude
+
+    return optimiser(
+        [
+            {"params": list(weight_decay_include)},
+            {"params": list(weight_decay_exclude), "weight_decay": 0.0},
+        ],
+        weight_decay=weight_decay,
+        lr=lr,
+    )

gradboard/scheduler.py

@@ -0,0 +1,214 @@
+"""
+Based on Smith (2017) https://arxiv.org/abs/1506.01186
+"""
+
+from typing import Optional
+import copy
+import math
+
+from torch.amp import GradScaler
+
+from .cycles import Cycle
+
+
+class PASS:
+    """
+    A self-configuring learning rate scheduler
+    """
+
+    def __init__(
+        self,
+        learning_rate_schedule: Cycle,
+        model,
+        optimiser,
+        scaler: Optional[GradScaler] = None,
+        range_test: bool = False,
+        cool_point_multiplier: float = 1 / 60,
+    ):
+        """
+        If not using range test, we assume the optimiser has the learning rates
+            set as desired.
+        """
+
+        self.model = model
+        self.optimiser = optimiser
+        self.scaler = scaler
+
+        self.learning_rate_schedule = learning_rate_schedule
+
+        self.range_test = range_test
+
+        self.original_param_groups = copy.deepcopy(optimiser.param_groups)
+
+        self.cool_point_multiplier = cool_point_multiplier
+
+        self.original_states = self._saved_states()
+
+        self.range_test_results = []
+
+        self.step_count = 0
+
+        if range_test:
+            self.start_range_test()  # sets LR to 1E-7
+
+    @property
+    def lr(self):
+        """
+        Return first lr from self.optimiser.param_groups
+            (this is used in learning rate range tests, in which case we can
+            assume they are all the same!)
+        """
+        for group in self.optimiser.param_groups:
+            return group["lr"]
+
+    @property
+    def in_range_test(self):
+        if not self.range_test:
+            return False
+        elif (len(self.range_test_results) == 0) or (
+            not math.isnan(self.range_test_results[-1][1])
+        ):
+            return True
+        else:
+            return False
+
+    @property
+    def trained(self):
+        if not self.range_test:
+            return True
+        elif math.isnan(self.range_test_results[-1][1]):
+            return True
+        else:
+            return False
+
+    @property
+    def finished(self):
+        return self.step_count >= len(self.learning_rate_schedule) - 1
+
+    def _saved_states(self):
+        saved_states = {
+            "model": copy.deepcopy(self.model.state_dict()),
+            "optimiser": copy.deepcopy(self.optimiser.state_dict()),
+        }
+        if self.scaler is not None:
+            saved_states["scaler"] = copy.deepcopy(self.scaler.state_dict())
+        return saved_states
+
+    def save_states(self):
+        self.saved_states = self._saved_states()
+
+    def load_states(self, saved_states):
+        self.model.load_state_dict(saved_states["model"])
+        self.optimiser.load_state_dict(saved_states["optimiser"])
+        if self.scaler is not None:
+            self.scaler.load_state_dict(saved_states["scaler"])
+
+    def recover_states(self):
+        self.load_states(self.saved_states)
+
+    @property
+    def _schedule_multiplier(self):
+        return self.learning_rate_schedule(
+            min(self.step_count, self.learning_rate_schedule.total_steps)
+        )
+
+    def set_all_lr(self, lr):
+        for group in self.optimiser.param_groups:
+            group["lr"] = lr
+
+    def start_range_test(self):
+        self.save_states()
+        self.optimiser.load_state_dict(self.original_states["optimiser"])
+        if self.scaler is not None:
+            self.scaler.load_state_dict(self.original_states["scaler"])
+        self.set_all_lr(1e-7)
+
+    def scale_all_lr(self, scaling_factor):
+        self.set_all_lr(self.lr * scaling_factor)
+
+    def end_range_test(self):
+        self.recover_states()
+        self.update_learning_rates()
+
+    def _smoothed_range_test(self, range_test_results):
+        range_test_results = sorted(range_test_results, key=lambda x: x[0])
+        learning_rates = [t[0] for t in range_test_results]
+        losses = [t[1] for t in self.range_test_results]
+        losses = losses[:-1] + [10 * max(losses)]
+        return list(zip(learning_rates, losses, strict=True))
+
+    def _plot_range_test(self, range_test_results):
+        """
+        Returns a tuple with x values (learning rates) and y values (losses)
+            which can then be passed to e.g. pyplot. We recommend presenting
+            the plot with a logarithmic x axis.
+        """
+        range_test_results = sorted(range_test_results, key=lambda x: x[0])
+        learning_rates = [t[0] for t in range_test_results]
+        losses = [t[1] for t in range_test_results]
+        return learning_rates, losses
+
+    def _apply_range_test_result(self):
+        """
+        ...
+        """
+        range_test_results = self._smoothed_range_test(self.range_test_results)
+        minimum = min(range_test_results, key=lambda x: x[1])
+        points_left_of_min = [r for r in range_test_results if r[0] < minimum[0]]
+        max_left_of_min = max(points_left_of_min, key=lambda x: x[1])
+        difference = max_left_of_min[1] - minimum[1]
+        max_lr = None
+        for p in sorted(points_left_of_min, key=lambda x: x[0]):
+            if (max_lr is None) and (p[1] < minimum[1] + 0.2 * difference):
+                max_lr = p[0]
+            else:
+                continue
+        self.set_all_lr(max_lr)
+        self.original_param_groups = copy.deepcopy(self.optimiser.param_groups)
+        print("High LR", max_lr)
+
+    def update_learning_rates(self):
+        if not self.finished:
+            for original, current in zip(
+                self.original_param_groups, self.optimiser.param_groups, strict=True
+            ):
+                base_lr = original["lr"]
+                min_lr = base_lr * self.cool_point_multiplier
+                current_lr = min_lr + (base_lr - min_lr) * self._schedule_multiplier
+                current["lr"] = current_lr
+
+    def _append_to_range_test(self, loss_item: float):
+
+        lr = self.lr
+
+        self.range_test_results.append((lr, loss_item))
+
+        if math.isnan(loss_item) or (lr >= 1.0):
+            self._apply_range_test_result()
+            self.end_range_test()
+        else:
+            # Continue range test, step up learning rate
+            self.scale_all_lr(1.05)
+
+    def step(self, loss_item: Optional[float] = None):
+        """
+        This function manages the process of
+            * Doing an initial range test
+            * Training for one microcycle using the learning rates from the
+                  initial range test ("burn in")
+            * Doing a second range test to set the learning rate schedule for
+                  the rest of training
+            * Updating learning rates during training according to the macrocycle
+        """
+        if self.in_range_test:  # True at init unless self.range_test = False
+            if not isinstance(loss_item, float):
+                raise ValueError(
+                    "When using range test functionality, "
+                    "`step()` expects a loss item."
+                )
+            self._append_to_range_test(loss_item)
+        elif self.trained and not self.finished:
+            self.step_count += 1
+            self.update_learning_rates()
+        else:
+            pass

gradboard-5.1.0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 nicholasbailey87
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

gradboard-5.1.0.dist-info/METADATA

@@ -0,0 +1,63 @@
+Metadata-Version: 2.3
+Name: gradboard
+Version: 5.1.0
+Summary: Easily snowboard down gnarly loss gradients
+License: MIT
+Author: Nicholas Bailey
+Requires-Python: >=3.8
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Description-Content-Type: text/markdown
+
+# gradboard
+![snowboarder](snowboarder.png "Image of a snowboarder")
+
+Easily snowboard down gnarly loss gradients
+
+## Getting started
+
+You can install gradboard with
+
+```
+pip install gradboard
+```
+
+PyTorch is a peer dependency of `gradboard`, which means
+  * You will need to make sure you have PyTorch installed in order to use `gradboard`
+  * PyTorch will **not** be installed automatically when you install `gradboard`
+
+We take this approach because PyTorch versioning is environment-specific and
+    we don't know where you will want to use `gradboard`. If we automatically install
+    PyTorch for you, there's a good chance we would get it wrong!
+
+Therefore, please also make sure you install PyTorch.
+
+## Usage examples
+
+### Decent model training outcomes without tuning hyperparameters
+
+`gradboard` includes
+
+  * An implementation of AdamS as proposed in Xie et al. (2023) "On the Overlooked
+        Pitfalls of Weight Decay and How to Mitigate Them: A Gradient-Norm
+        Perspective" (https://openreview.net/pdf?id=vnGcubtzR1), which in practice
+        makes model training more robust to the weight decay setting.
+  * Utilities for implementing popular learning rate schedules
+  * An implementation of an automatic max/min learning rate finder based on Smith
+        (2017) "Cyclical Learning Rates for Training Neural Networks"
+        (https://arxiv.org/abs/1506.01186)
+  * Sensible defaults
+
+In practice this means that you can train a neural network and get decent performance
+    right out of the box, just by using the `PASS` (point-and-shoot scheduler), even
+    for unfamiliar architectures or problem domains.
+
+
+
+

gradboard-5.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+gradboard/__init__.py,sha256=57AkHusYwLCsusiVnajH5pMFKioRCj-3IjF9qpdOzE0,69
+gradboard/cycles.py,sha256=XnXNzCBI3J7OmzZQ3bKItffeDVDMXhTK6qicuYow6_4,10507
+gradboard/optimiser.py,sha256=ds7rk67eiOYRsSwJyAWp5nK2vhCtko03FU9vHPUSC6E,1417
+gradboard/scheduler.py,sha256=u3ojuGJQ4ZVn5fZ9L49CsFTu0TS-iqijarSEkmIWIfA,7056
+gradboard-5.1.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+gradboard-5.1.0.dist-info/METADATA,sha256=b0TWoWtSAJj2RTWleIHD_SquassUmdfR1rvIyXXM-Fo,2246
+gradboard-5.1.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+gradboard-5.1.0.dist-info/RECORD,,

gradboard-5.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.1.3
+Root-Is-Purelib: true
+Tag: py3-none-any