diff --git a/docs/full/lib/model.rst b/docs/full/lib/model.rst index 84e8d49e33..37213553b6 100755 --- a/docs/full/lib/model.rst +++ b/docs/full/lib/model.rst @@ -33,6 +33,22 @@ networks package :include-all-objects: :noindex: +optimizers package +================== + +.. automodapi:: lib.model.optimizers.adabelief + :include-all-objects: + :noindex: + +| +.. automodapi:: lib.model.optimizers.lion + :include-all-objects: + :noindex: + +| +.. automodapi:: lib.model.optimizers.keras_legacy + :include-all-objects: + :noindex: model package ============= @@ -62,7 +78,3 @@ model package | .. automodapi:: lib.model.normalization :include-all-objects: - -| -.. automodapi:: lib.model.optimizers - :include-all-objects: diff --git a/docs/full/lib/training.rst b/docs/full/lib/training.rst index cbe7382d34..f6609db92c 100644 --- a/docs/full/lib/training.rst +++ b/docs/full/lib/training.rst @@ -20,6 +20,11 @@ The training Package handles libraries to assist with training a model :include-all-objects: :no-inheritance-diagram: +| +.. automodapi:: lib.training.optimizer + :include-all-objects: + :no-inheritance-diagram: + | .. automodapi:: lib.training.preview :include-all-objects: diff --git a/lib/align/aligned_mask.py b/lib/align/aligned_mask.py index 28a95976c8..33c38416d9 100644 --- a/lib/align/aligned_mask.py +++ b/lib/align/aligned_mask.py @@ -466,7 +466,7 @@ def __init__(self, blur_kernel: int = 0, blur_type: T.Literal["gaussian", "normalized"] | None = "gaussian", blur_passes: int = 1) -> None: - logger.debug(parse_class_init(locals())) + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] self._area = area self._landmark_type = landmark_type self._landmarks = landmarks diff --git a/lib/config/objects.py b/lib/config/objects.py index d496c56904..c2800e52b0 100644 --- a/lib/config/objects.py +++ b/lib/config/objects.py @@ -1,5 +1,5 @@ #! /usr/env/bin/python3 -""" Dataclass objects for holding and validating Faceswap Config items """ +"""Dataclass objects for holding and validating Faceswap Config item""" from __future__ import annotations import gettext @@ -26,7 +26,7 @@ # TODO allow list items other than strings @dataclass class ConfigItem(Generic[T]): # pylint:disable=too-many-instance-attributes - """ A dataclass for storing config items loaded from config.ini files and dynamically assigning + """A dataclass for storing config items loaded from config.ini files and dynamically assigning and validating that the correct datatype is used. The value loaded from the .ini config file can be accessed with either: @@ -37,17 +37,17 @@ class ConfigItem(Generic[T]): # pylint:disable=too-many-instance-attributes Parameters ---------- - datatype : type + datatype A python type class. This limits the type of data that can be provided in the .ini file and ensures that the value is returned to faceswap is correct. Valid datatypes are: `int`, `float`, `str`, `bool` or `list`. Note that `list` items must all be strings. - default : Any + default The default value for this option. It must be of the same type as :attr:`datatype`. - group : str + group The group that this config item exists within in the config section - info : str + info A description of what this option does. - choices : list[str] | Literal["colorchooser"], optional + choices If this option's datatype is a `str` then valid selections can be defined here, empty list for any value. If the option's datatype is a `list`, then this option must be populated with the valid selections. This validates the option and also enables a combobox / radio @@ -55,61 +55,60 @@ class ConfigItem(Generic[T]): # pylint:disable=too-many-instance-attributes literal "colorchooser" to present a color choosing interface in the GUI. Ignored for all other datatypes Default: [] (empty list: no options) - gui_radio : bool, optional + gui_radio If :attr:`choices` are defined, this indicates that the GUI should use radio buttons rather than a combobox to display this option. Default: ``False`` - min_max : tuple[int | float, int | float] | None, optional + min_max For `int` and `float` :attr:`datatype` this is required otherwise it is ignored. Should be a tuple of min and max accepted values of the same datatype as the option value. This is used for controlling the GUI slider range. Values are not enforced. Default: ``None`` - rounding : int | None, optional + rounding For `int` and `float :attr:datatypes this is required to be > 0 otherwise it is ignored. Used for the GUI slider. For `float`, this is the number of decimal places to display. For `int` this is the step size. Default: `-1` (ignored) - fixed : bool, optional + fixed [train only]. Training configurations are fixed when the model is created, and then reloaded from the state file. Marking an item as fixed=``False`` indicates that this value can be changed for existing models, and will override the value saved in the state file with the updated value in config. Default: ``True`` """ datatype: type[T] - """ type : A python type class. The datatype of the config value. One of `int`, `float`, `str`, - `bool` or `list`. `list` will only contain `str` items """ + """A python type class. The datatype of the config value. One of `int`, `float`, `str`, `bool` + or `list`. `list` will only contain `str` items""" default: T - """ Any : The default value for this option. It is of the same type as :attr:`datatype` """ + """The default value for this option. It is of the same type as :attr:`datatype`""" group: str - """ str : The group that this config option belongs to """ + """The group that this config option belongs to""" info: str - """ str : A description of what this option does """ + """A description of what this option does""" choices: list[str] | Literal["colorchooser"] = field(default_factory=list) - """ list[str] | Literal["colorchooser"]: If this option's datatype is a `str` then valid - selections may be defined here, Empty list if any value is valid. If the datatype is a `list` - then valid choices will be populated here. If the default value is a hex color code, then the - literal "colorchooser" will display a color choosing interface in the GUI. """ + """If this option's datatype is a `str` then valid selections may be defined here, Empty list + if any value is valid. If the datatype is a `list` then valid choices will be populated here. + If the default value is a hex color code, then the literal "colorchooser" will display a color + choosing interface in the GUI.""" gui_radio: bool = False - """ bool : indicates that the GUI should use radio buttons rather than a combobox to display - this option if :attr:`choices` is populated """ + """indicates that the GUI should use radio buttons rather than a combobox to display this + option if :attr:`choices` is populated""" min_max: tuple[T, T] | None = None - """ tuple[int | float, int | float] | None : For `int` and `float` :attr:`datatype` this will - be populated otherwise it will be ``None``. Used for controlling the GUI slider range. Values - are not enforced. """ + """For `int` and `float` :attr:`datatype` this will be populated otherwise it will be ``None``. + Used for controlling the GUI slider range. Values are not enforced.""" rounding: int = -1 - """ int : For `int` and `float` :attr:`datatypes` this will be > 0 otherwise it will be `-1`. - Used for the GUI slider. For `float`, this is the number of decimal places to display. For - `int` this is the step size. """ + """For `int` and `float` :attr:`datatypes` this will be > 0 otherwise it will be `-1`. Used for + the GUI slider. For `float`, this is the number of decimal places to display. For `int` this is + the step size.""" fixed: bool = True - """ bool : Only used for train.model configurations. Options marked as fixed=``False`` - indicates that this value can be changed for existing models, otherwise the option set when the - model commenced training is fixed and cannot be changed. Default: ``True`` """ + """Only used for train.model configurations. Options marked as fixed=``False`` indicates that + this value can be changed for existing models, otherwise the option set when the model + commenced training is fixed and cannot be changed. Default: ``True``""" _value: T = field(init=False) - """ Any : The value of the config item of type :attr:`datatype`""" + """The value of the config item of type :attr:`datatype`""" _name: str = field(init=False) - """ str: The option name for this object. Set when the config is first loaded """ + """The option name for this object. Set when the config is first loaded""" @property def helptext(self) -> str: - """ str | Description of the config option with additional formating and helptext added - from the item parameters """ + """Description of the config option with additional formatting and helptext added from the + item parameters""" retval = f"{self.info}\n" if not self.fixed: retval += _("\nThis option can be updated for existing models.\n") @@ -122,20 +121,20 @@ def helptext(self) -> str: retval += _("\nChoose from: True, False") elif self.datatype == int: assert self.min_max is not None - cmin, cmax = self.min_max - retval += _("\nSelect an integer between {} and {}").format(cmin, cmax) + c_min, c_max = self.min_max + retval += _("\nSelect an integer between {} and {}").format(c_min, c_max) elif self.datatype == float: assert self.min_max is not None - cmin, cmax = self.min_max - retval += _("\nSelect a decimal number between {} and {}").format(cmin, cmax) + c_min, c_max = self.min_max + retval += _("\nSelect a decimal number between {} and {}").format(c_min, c_max) default = ", ".join(self.default) if isinstance(self.default, list) else self.default retval += _("\n[Default: {}]").format(default) return retval @property def value(self) -> T: - """ Any : The config value for this item loaded from the config .ini file. String values - will always be lowercase, regardless of what is loaded from Config """ + """The config value for this item loaded from the config .ini file. String values will + always be lowercase, regardless of what is loaded from Config""" retval = self._value if isinstance(self._value, str): retval = cast(T, self._value.lower()) @@ -145,35 +144,34 @@ def value(self) -> T: @property def ini_value(self) -> str: - """ str : The current value of the ConfigItem as a string for writing to a .ini file """ + """The current value of the ConfigItem as a string for writing to a .ini file""" if isinstance(self._value, list): return ", ".join(str(x) for x in self._value) return str(self._value) @property def name(self) -> str: - """str: The name associated with this option """ + """The name associated with this option""" return self._name def _validate_type(self, # pylint:disable=too-many-return-statements expected_type: Any, attr: Any, depth=1) -> bool: - """ Validate that provided types are correct when this Dataclass is initialized + """Validate that provided types are correct when this Dataclass is initialized Parameters ---------- - expected_type : Any + expected_type The expected data type for the given attribute - attr : Any + attr The attribute to test for correctness - depth : int, optional + depth The current recursion depth Returns ------- - bool - ``True`` if the given attribute is a valid datatype + ``True`` if the given attribute is a valid datatype Raises ------ @@ -218,7 +216,7 @@ def _validate_type(self, # pylint:disable=too-many-return-statements return False def _validate_required(self) -> None: - """ Validate that required parameters are populated + """Validate that required parameters are populated Raises ------ @@ -231,7 +229,7 @@ def _validate_required(self) -> None: raise ValueError("Option info must me provided") def _validate_choices(self) -> None: - """ Validate that choices have been used correctly + """Validate that choices have been used correctly Raises ------ @@ -266,7 +264,7 @@ def _validate_choices(self) -> None: raise ValueError("Config item of type list must have choices defined") def _validate_numeric(self) -> None: - """ Validate that float and int values have been set correctly + """Validate that float and int values have been set correctly Raises ------ @@ -283,7 +281,7 @@ def _validate_numeric(self) -> None: f") values. Got {self.min_max}") def __post_init__(self) -> None: - """ Validate and type check that the given parameters are valid and set the default value. + """Validate and type check that the given parameters are valid and set the default value. Raises ------ @@ -303,27 +301,25 @@ def __post_init__(self) -> None: self._validate_numeric() def get(self) -> T: - """ Obtain the currently stored configuration value + """Obtain the currently stored configuration value Returns ------- - Any - The config value for this item loaded from the config .ini file. String values will - always be lowecase, regardless of what is loaded from Config """ + The config value for this item loaded from the config .ini file. String values will always + be lowercase, regardless of what is loaded from Config""" return self.value def _parse_list(self, value: str | list[str]) -> list[str]: - """ Parse inbound list values. These can be space/comma-separated strings or a list. + """Parse inbound list values. These can be space/comma-separated strings or a list. Parameters ---------- - value : str | list[str] + value The inbound value to be converted to a list Returns ------- - list[str] - List of strings representing the inbound values. + List of strings representing the inbound values. """ if not value: return [] @@ -335,17 +331,15 @@ def _parse_list(self, value: str | list[str]) -> list[str]: return retval def _validate_selection(self, value: str | list[str]) -> str | list[str]: - """ Validate that the given value is valid within the stored choices + """Validate that the given value is valid within the stored choices Parameters ---------- - str | list[str] - The inbound config value to validate + The inbound config value to validate Returns ------- - bool - ``True`` if the selected value is a valid choice + ``True`` if the selected value is a valid choice """ assert isinstance(self.choices, list) choices = [x.lower() for x in self.choices] @@ -370,11 +364,11 @@ def _validate_selection(self, value: str | list[str]) -> str | list[str]: return valid def set(self, value: T) -> None: - """ Set the item's option value + """Set the item's option value Parameters ---------- - value : Any + value The value to set this item to. Must be of type :attr:`datatype` Raises @@ -409,11 +403,11 @@ def set(self, value: T) -> None: self._value = value def set_name(self, name: str) -> None: - """ Set the logging name for this object for display purposes + """Set the logging name for this object for display purposes Parameters ---------- - name : str + name The name to assign to this option """ logger.debug("Setting name to '%s'", name) @@ -421,40 +415,48 @@ def set_name(self, name: str) -> None: self._name = name def __call__(self) -> T: - """ Obtain the currently stored configuration value + """Obtain the currently stored configuration value Returns ------- - Any - The config value for this item loaded from the config .ini file. String values will - always be lowecase, regardless of what is loaded from Config """ + The config value for this item loaded from the config .ini file. String values will always + be lowercase, regardless of what is loaded from Config""" return self.value @dataclass class ConfigSection: - """ Dataclass for holding information about configuration sections and the contained + """Dataclass for holding information about configuration sections and the contained configuration items Parameters ---------- - helptext : str + helptext The helptext to be displayed for the configuration section - options : dict[str, :class:`ConfigItem`] + options Dictionary of configuration option name to the options for the section """ helptext: str options: dict[str, ConfigItem] +class ConfigReprMeta(type): + """A custom repr for printing currently selected config values""" + def __repr__(cls) -> str: + params = ", ".join(f"{k}={repr(v.value)}" + for k, v in cls.__dict__.items() + if isinstance(v, ConfigItem)) + return f"{cls.__name__}({params})" + + @dataclass -class GlobalSection: - """ A dataclass for holding and identifying global sub-sections for plugin groups. Any global +class GlobalSection(metaclass=ConfigReprMeta): + """A dataclass for holding and identifying global sub-sections for plugin groups. Any global subsections must inherit from this. Parameters ---------- - helptext : str + helptext The helptext to be displayed for the global configuration section """ helptext: str diff --git a/lib/model/autoclip.py b/lib/model/autoclip.py index 03d1a54af7..384418d41c 100644 --- a/lib/model/autoclip.py +++ b/lib/model/autoclip.py @@ -1,29 +1,28 @@ -""" Auto clipper for clipping gradients. """ +"""Auto clipper for clipping gradients.""" from __future__ import annotations import logging -import typing as T +import math +from collections import deque import numpy as np import torch +from torch import nn from lib.logger import parse_class_init from lib.utils import get_module_objects -if T.TYPE_CHECKING: - from keras import KerasTensor - logger = logging.getLogger(__name__) class AutoClipper(): - """ AutoClip: Adaptive Gradient Clipping for Source Separation Networks + """AutoClip: Adaptive Gradient Clipping for Source Separation Networks Parameters ---------- - clip_percentile: int + clip_percentile The percentile to clip the gradients at - history_size: int, optional + history_size The number of iterations of data to use to calculate the norm Default: ``10000`` References @@ -33,32 +32,32 @@ class AutoClipper(): """ def __init__(self, clip_percentile: int, history_size: int = 10000) -> None: logger.debug(parse_class_init(locals())) - self._clip_percentile = clip_percentile - self._history_size = history_size - self._grad_history: list[float] = [] + self._grad_history: deque[float] = deque(maxlen=history_size) - logger.debug("Initialized %s", self.__class__.__name__) - - def __call__(self, gradients: list[KerasTensor]) -> list[KerasTensor]: - """ Call the AutoClip function. + def __call__(self, parameters: list[nn.Parameter], *args) -> None: + """Call the AutoClip function. Parameters ---------- - gradients: list[:class:`keras.KerasTensor`] - The list of gradient tensors for the optimizer - - Returns - ---------- - list[:class:`keras.KerasTensor`] - The autoclipped gradients + parameters + The parameters to clip + args + Unused but for compatibility """ - self._grad_history.append(sum(g.data.norm(2).item() ** 2 - for g in gradients if g is not None) ** (1. / 2)) - self._grad_history = self._grad_history[-self._history_size:] - clip_value = np.percentile(self._grad_history, self._clip_percentile) - torch.nn.utils.clip_grad_norm_(gradients, T.cast(float, clip_value)) - return gradients + with torch.no_grad(): + norms = [p.grad.norm(2).item() for p in parameters if p.grad is not None] + + if not norms: + return + + global_norm = sum(n ** 2 for n in norms) ** 0.5 + if not math.isfinite(global_norm): + return + + self._grad_history.append(global_norm) + clip_value = float(np.percentile(self._grad_history, self._clip_percentile)) + nn.utils.clip_grad_norm_(parameters, clip_value) __all__ = get_module_objects(__name__) diff --git a/lib/model/losses/feature_loss.py b/lib/model/losses/feature_loss.py index a9c1e9092d..c127a47b97 100644 --- a/lib/model/losses/feature_loss.py +++ b/lib/model/losses/feature_loss.py @@ -133,9 +133,9 @@ def _normalize_output(cls, inputs: torch.Tensor, epsilon: float = 1e-10) -> torc Parameters ---------- - inputs: :class:`keras.KerasTensor` + inputs An output tensor from the trunk model - epsilon: float, optional + epsilon Epsilon to apply to the normalization operation. Default: `1e-10` """ norm_factor = torch.sqrt(torch.sum(torch.square(inputs), dim=1, keepdim=True)) diff --git a/lib/model/optimizers/__init__.py b/lib/model/optimizers/__init__.py new file mode 100644 index 0000000000..fd5ddd0642 --- /dev/null +++ b/lib/model/optimizers/__init__.py @@ -0,0 +1,5 @@ +#! /usr/env/bin/python3 +"""Custom Torch Optimizers""" +from .adabelief import AdaBelief +from .lion import Lion +from .keras_legacy import AdaBelief as AdaBeliefKeras diff --git a/lib/model/optimizers/adabelief.py b/lib/model/optimizers/adabelief.py new file mode 100644 index 0000000000..9d85a46a83 --- /dev/null +++ b/lib/model/optimizers/adabelief.py @@ -0,0 +1,287 @@ +#! /usr/env/bin/python3 +"""AdaBelief optimizer for Torch""" +# BSD 2-Clause License +# +# Copyright (c) 2021, Juntang Zhuang +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import logging +import math +import typing as T + +import torch +from torch.optim.optimizer import Optimizer + +from lib.logger import parse_class_init +from lib.utils import get_module_objects + +logger = logging.getLogger(__name__) + + +class AdaBelief(Optimizer): + """Implements AdaBelief algorithm. Modified from Adam in PyTorch + + Parameters + ---------- + params + Iterable of parameters to optimize or dicts defining parameter groups + lr + Learning rate. Default: 1e-3 + betas + Coefficients used for computing running averages of gradient and its square. + Default: (0.9, 0.999) + eps + Term added to the denominator to improve numerical stability. Default: 1e-16 + weight_decay + Weight decay (L2 penalty). Default: 0 + amsgrad + Whether to use the AMSGrad variant of this algorithm from the paper `On the Convergence + of Adam and Beyond`. Default: ``False`` + weight_decouple + If set as True, then the optimizer uses decoupled weight decay as in AdamW. + Default: ``True`` + fixed_decay + This is used when weight_decouple is set as True. + - When fixed_decay == True, the weight decay is performed as W_{new} = W_{old} - W_{old} + * decay. + - When fixed_decay == False, the weight decay is performed as W_{new} = W_{old} - W_{old} + * decay * lr. Note that in this case, the weight decay ratio decreases with learning rate + (lr). + Default: ``False`` + rectify + If set as True, then perform the rectified update similar to RAdam. + Default: ``True`` + degenerated_to_sgd + If set as True, then perform SGD update when variance of gradient is high. + Default: ``True`` + + Reference + --------- + AdaBelief Optimizer, adapting step sizes by the belief in observed gradients, NeurIPS 2020 + https://github.com/juntang-zhuang/Adabelief-Optimizer + """ + def __init__(self, # pylint:disable=too-many-positional-arguments,too-many-arguments # noqa[C901] + params: T.Iterable, + lr: float = 1e-3, + betas: tuple[float, float] = (0.9, 0.999), + eps: float = 1e-16, + weight_decay: float = 0.0, + amsgrad: bool = False, + weight_decouple: bool = True, + fixed_decay: bool = False, + rectify: bool = True, + degenerated_to_sgd: bool = True) -> None: + logger.debug(parse_class_init(locals())) + if 0.0 > lr: + raise ValueError(f"Invalid learning rate: {lr}") + if 0.0 > eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") + + self.degenerated_to_sgd = degenerated_to_sgd + if isinstance(params, (list, tuple)) and len(params) > 0 and isinstance(params[0], + dict): + for param in params: + if "betas" in param and (param["betas"][0] != betas[0] + or param["betas"][1] != betas[1]): + param["buffer"] = [[None, None, None] for _ in range(10)] + + defaults = {"lr": lr, + "betas": betas, + "eps": eps, + "weight_decay": weight_decay, + "amsgrad": amsgrad, + "buffer": [[None, None, None] for _ in range(10)]} + super().__init__(params, defaults) + + self.degenerated_to_sgd = degenerated_to_sgd + self.weight_decouple = weight_decouple + self.rectify = rectify + self.fixed_decay = fixed_decay + if self.weight_decouple: + logger.debug("[AdaBelief] Weight decoupling enabled in AdaBelief") + if self.fixed_decay: + logger.debug("[AdaBelief] Weight decay fixed") + if self.rectify: + logger.debug("[AdaBelief] Rectification enabled in AdaBelief") + if amsgrad: + logger.debug("[AdaBelief] AMSGrad enabled in AdaBelief") + + def __setstate__(self, state: dict[str, T.Any]) -> None: + """Set parameter state""" + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("amsgrad", False) + + def reset(self) -> None: + """Reset parameters""" + for group in self.param_groups: + for p in group["params"]: + state = self.state[p] + amsgrad = group["amsgrad"] + + # State initialization + state["step"] = torch.zeros((), dtype=torch.float32) + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like(p.data, memory_format=torch.preserve_format) + + # Exponential moving average of squared gradient values + state["exp_avg_var"] = torch.zeros_like(p.data, + memory_format=torch.preserve_format) + + if amsgrad: + # Maintains max of all exp. moving avg. of sq. grad. values + state["max_exp_avg_var"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format) + + def step(self, # type:ignore[override] # noqa[C901] + closure: T.Callable | None = None) -> torch.Tensor: + """Performs a single optimization step. + + Parameters + ---------- + closure + A closure that reevaluates the model and returns the loss. Default: ``None`` + """ + # pylint:disable=duplicate-code,too-many-statements,too-many-branches,too-many-locals + loss: torch.Tensor | None = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + + # cast data type + half_precision = False + if p.data.dtype == torch.float16: + half_precision = True + p.data = p.data.float() + p.grad = p.grad.float() + + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError( + "AdaBelief does not support sparse gradients, please consider SparseAdam " + "instead") + amsgrad = group["amsgrad"] + + state = self.state[p] + + beta1, beta2 = group["betas"] + + # State initialization + if len(state) == 0: + state["step"] = torch.zeros((), dtype=torch.float32) + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like(p.data, + memory_format=torch.preserve_format) + # Exponential moving average of squared gradient values + state["exp_avg_var"] = torch.zeros_like(p.data, + memory_format=torch.preserve_format) + if amsgrad: + # Maintains max of all exp. moving avg. of sq. grad. values + state["max_exp_avg_var"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format) + + # perform weight decay, check if decoupled weight decay + if self.weight_decouple: + if not self.fixed_decay: + p.data.mul_(1.0 - group["lr"] * group["weight_decay"]) + else: + p.data.mul_(1.0 - group["weight_decay"]) + else: + if group["weight_decay"] != 0: + grad.add_(p.data, alpha=group["weight_decay"]) + + # get current state variable + exp_avg, exp_avg_var = state["exp_avg"], state["exp_avg_var"] + + state["step"] += 1 + bias_correction1 = 1 - beta1 ** state["step"] + bias_correction2 = 1 - beta2 ** state["step"] + + # Update first and second moment running average + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + grad_residual = grad - exp_avg + exp_avg_var.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1 - beta2) + + if amsgrad: + max_exp_avg_var = state["max_exp_avg_var"] + # Maintains the maximum of all 2nd moment running avg. till now + torch.max(max_exp_avg_var, exp_avg_var.add_(group["eps"]), out=max_exp_avg_var) + + # Use the max. for normalizing running avg. of gradient + denom = (max_exp_avg_var.sqrt() / + math.sqrt(bias_correction2)).add_(group["eps"]) + else: + denom = (exp_avg_var.add_(group["eps"]).sqrt() / + math.sqrt(bias_correction2)).add_(group["eps"]) + + # update + if not self.rectify: + # Default update + step_size = group["lr"] / bias_correction1 + p.data.addcdiv_(exp_avg, denom, value=-step_size) + + else: # Rectified update, forked from RAdam + buffered = group["buffer"][int(state["step"] % 10)] + if state["step"] == buffered[0]: + n_sma, step_size = buffered[1], buffered[2] + else: + buffered[0] = state["step"] + beta2_t = beta2 ** state["step"] + n_sma_max = 2 / (1 - beta2) - 1 + n_sma = n_sma_max - 2 * state["step"] * beta2_t / (1 - beta2_t) + buffered[1] = n_sma + + # more conservative since it"s an approximated value + if n_sma >= 5: + step_size = math.sqrt( + (1 - beta2_t) * (n_sma - 4) / + (n_sma_max - 4) * (n_sma - 2) / + n_sma * n_sma_max / (n_sma_max - 2)) / (1 - beta1 ** state["step"]) + elif self.degenerated_to_sgd: + step_size = 1.0 / (1 - beta1 ** state["step"]) + else: + step_size = -1 + buffered[2] = step_size + + if n_sma >= 5: + denom = exp_avg_var.sqrt().add_(group["eps"]) + p.data.addcdiv_(exp_avg, denom, value=-step_size * group["lr"]) + elif step_size > 0: + p.data.add_(exp_avg, alpha=-step_size * group["lr"]) + + if half_precision: + p.data = p.data.half() + p.grad = p.grad.half() + + return T.cast(torch.Tensor, loss) + + +__all__ = get_module_objects(__name__) diff --git a/lib/model/optimizers.py b/lib/model/optimizers/keras_legacy.py similarity index 76% rename from lib/model/optimizers.py rename to lib/model/optimizers/keras_legacy.py index 835258ad28..530b36ad37 100644 --- a/lib/model/optimizers.py +++ b/lib/model/optimizers/keras_legacy.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -""" Custom Optimizers for Torch/keras """ +"""Legacy keras Optimizers for weight migration""" from __future__ import annotations import inspect import logging @@ -12,13 +12,14 @@ from lib.utils import get_module_objects if T.TYPE_CHECKING: - from keras import KerasTensor, Variable + from torch import Tensor + from keras import Variable logger = logging.getLogger(__name__) class AdaBelief(Optimizer): # pylint:disable=too-many-instance-attributes,too-many-ancestors - """ Implementation of the AdaBelief Optimizer + """Implementation of the AdaBelief Optimizer Inherits from: keras.optimizers.Optimizer. @@ -32,30 +33,30 @@ class AdaBelief(Optimizer): # pylint:disable=too-many-instance-attributes,too-m Parameters ---------- - learning_rate: `Tensor`, float or :class: `keras.optimizers.schedules.LearningRateSchedule` + learning_rate The learning rate. - beta_1: float + beta_1 The exponential decay rate for the 1st moment estimates. - beta_2: float + beta_2 The exponential decay rate for the 2nd moment estimates. - epsilon: float + epsilon A small constant for numerical stability. - amsgrad: bool + amsgrad Whether to apply AMSGrad variant of this algorithm from the paper "On the Convergence of Adam and beyond". - rectify: bool + rectify Whether to enable rectification as in RectifiedAdam - sma_threshold. float + sma_threshold The threshold for simple mean average. - total_steps: int + total_steps Total number of training steps. Enable warmup by setting a positive value. - warmup_proportion: float + warmup_proportion The proportion of increasing steps. - min_lr: float + min_lr Minimum learning rate after warmup. - name: str, optional + name Name for the operations created when applying gradients. Default: ``"AdaBeliefOptimizer"``. - **kwargs: dict + **kwargs Standard Keras Optimizer keyword arguments. Allowed to be (`weight_decay`, `clipnorm`, `clipvalue`, `global_clipnorm`, `use_ema`, `ema_momentum`, `ema_overwrite_frequency`, `loss_scale_factor`, `gradient_accumulation_steps`) @@ -90,7 +91,7 @@ class AdaBelief(Optimizer): # pylint:disable=too-many-instance-attributes,too-m References ---------- - Juntang Zhuang et al. - AdaBelief Optimizer: Adapting stepsizes by the belief in observed + Juntang Zhuang et al. - AdaBelief Optimizer: Adapting step sizes by the belief in observed gradients - https://arxiv.org/abs/2010.07468. Original implementation - https://github.com/juntang-zhuang/Adabelief-Optimizer @@ -148,13 +149,9 @@ def __init__(self, # pylint:disable=too-many-arguments,too-many-positional-argu self.amsgrad = amsgrad self.rectify = rectify self.sma_threshold = sma_threshold - # TODO change the following 2 to "warm_up_steps" - # TODO Make learning rate warm up a global option - # Or these params can be calculated from a user "warm_up_steps" parameter self.total_steps = total_steps self.warmup_proportion = warmup_proportion self.min_learning_rate = min_learning_rate - logger.debug("Initialized %s", self.__class__.__name__) self._momentums: list[Variable] = [] self._velocities: list[Variable] = [] @@ -168,7 +165,7 @@ def build(self, variables: list[Variable]) -> None: Parameters ---------- - variables: list[:class:`keras.Variable`] + variables list of model variables to build AdaBelief variables on. """ if self.built: @@ -187,20 +184,19 @@ def build(self, variables: list[Variable]) -> None: logger.debug("Built AdaBelief. momentums: %s, velocities: %s, velocity_hats: %s", len(self._momentums), len(self._velocities), len(self._velocity_hats)) - def _maybe_warmup(self, learning_rate: KerasTensor, local_step: KerasTensor) -> KerasTensor: - """ Do learning rate warm up if requested + def _maybe_warmup(self, learning_rate: Tensor, local_step: Tensor) -> Tensor: + """Do learning rate warm up if requested Parameters ---------- - learning_rate: :class:`keras.KerasTensor` + learning_rate The learning rate - local_step: :class:`keras.KerasTensor` + local_step The current training step Returns ------- - :class:`keras.KerasTensor` - Either the original learning rate or adjusted learning rate if warmup is requested + Either the original learning rate or adjusted learning rate if warmup is requested """ if self.total_steps <= 0: return learning_rate @@ -210,74 +206,78 @@ def _maybe_warmup(self, learning_rate: KerasTensor, local_step: KerasTensor) -> min_lr = ops.cast(self.min_learning_rate, learning_rate.dtype) decay_steps = ops.maximum(total_steps - warmup_steps, 1) decay_rate = ops.divide(min_lr - learning_rate, decay_steps) - return ops.where(local_step <= warmup_steps, - ops.multiply(learning_rate, (ops.divide(local_step, warmup_steps))), - ops.multiply(learning_rate + decay_rate, - ops.minimum(local_step - warmup_steps, decay_steps))) + return T.cast("Tensor", + ops.where(local_step <= warmup_steps, + ops.multiply(learning_rate, + (ops.divide(local_step, warmup_steps))), + ops.multiply(learning_rate + decay_rate, + ops.minimum(local_step - warmup_steps, decay_steps)))) def _maybe_rectify(self, - momentum: KerasTensor, - velocity: KerasTensor, - local_step: KerasTensor, - beta_2_power: KerasTensor) -> KerasTensor: - """ Apply rectification, if requested + momentum: Tensor, + velocity: Tensor, + local_step: Tensor, + beta_2_power: Tensor) -> Tensor: + """Apply rectification, if requested Parameters ---------- - momentum: :class:`keras.KerasTensor` + momentum The momentum update - velocity: :class:`keras.KerasTensor` + velocity The velocity update - local_step: :class:`keras.KerasTensor` + local_step The current training step beta_2_power Adjusted exponential decay rate for the 2nd moment estimates. Returns ------- - :class:`keras.KerasTensor` - The standard or rectified update (if rectification enabled) + The standard or rectified update (if rectification enabled) """ if not self.rectify: - return ops.divide(momentum, ops.add(velocity, self.epsilon)) + return T.cast("Tensor", ops.divide(momentum, ops.add(velocity, self.epsilon))) sma_inf = 2 / (1 - self.beta_2) - 1 sma_t = sma_inf - 2 * local_step * beta_2_power / (1 - beta_2_power) rect = ops.sqrt((sma_t - 4) / (sma_inf - 4) * (sma_t - 2) / (sma_inf - 2) * sma_inf / sma_t) - return ops.where(sma_t >= self.sma_threshold, - ops.divide( - ops.multiply(rect, momentum), - (ops.add(velocity, self.epsilon))), - momentum) + return T.cast("Tensor", + ops.where(sma_t >= self.sma_threshold, + ops.divide(ops.multiply(rect, momentum), + (ops.add(velocity, self.epsilon))), + momentum)) def update_step(self, - gradient: KerasTensor, + gradient: Tensor, variable: Variable, - learning_rate: Variable) -> None: + learning_rate: Tensor) -> None: """Update step given gradient and the associated model variable for AdaBelief. Parameters ---------- - gradient :class:`keras.KerasTensor` + gradient The gradient to update - variable: :class:`keras.Variable` + variable The variable to update - learning_rate: :class:`keras.Variable` + learning_rate The learning rate """ - local_step = ops.cast(self.iterations + 1, variable.dtype) - learning_rate = self._maybe_warmup(ops.cast(learning_rate, variable.dtype), local_step) - gradient = ops.cast(gradient, variable.dtype) + local_step = T.cast("Tensor", ops.cast(self.iterations + 1, variable.dtype)) + learning_rate = self._maybe_warmup(T.cast("Tensor", + ops.cast(learning_rate, variable.dtype)), + local_step) + gradient = T.cast("Tensor", ops.cast(gradient, variable.dtype)) beta_1_power = ops.power(ops.cast(self.beta_1, variable.dtype), local_step) - beta_2_power = ops.power(ops.cast(self.beta_2, variable.dtype), local_step) + beta_2_power = T.cast("Tensor", + ops.power(ops.cast(self.beta_2, variable.dtype), local_step)) # m_t = b1 * m + (1 - b1) * g # => m_t = m + (g - m) * (1 - b1) - momentum = self._momentums[self._get_variable_index(variable)] + momentum = T.cast("Variable", self._momentums[self._get_variable_index(variable)]) self.assign_add(momentum, ops.multiply(ops.subtract(gradient, momentum), 1 - self.beta_1)) - momentum_corr = ops.divide(momentum, (1 - beta_1_power)) + momentum_corr = T.cast("Tensor", ops.divide(momentum, (1 - beta_1_power))) # v_t = b2 * v + (1 - b2) * (g - m_t)^2 + e # => v_t = v + ((g - m_t)^2 - v) * (1 - b2) + e @@ -291,16 +291,17 @@ def update_step(self, if self.amsgrad: velocity_hat = self._velocity_hats[self._get_variable_index(variable)] self.assign(velocity_hat, ops.maximum(velocity, velocity_hat)) - velocity_corr = ops.sqrt(ops.divide(velocity_hat, (1 - beta_2_power))) + velocity_corr = T.cast("Tensor", + ops.sqrt(ops.divide(velocity_hat, (1 - beta_2_power)))) else: - velocity_corr = ops.sqrt(ops.divide(velocity, (1 - beta_2_power))) + velocity_corr = T.cast("Tensor", ops.sqrt(ops.divide(velocity, (1 - beta_2_power)))) var_t = self._maybe_rectify(momentum_corr, velocity_corr, local_step, beta_2_power) self.assign_sub(variable, ops.multiply(learning_rate, var_t)) def get_config(self) -> dict[str, T.Any]: - """ Returns the config of the optimizer. + """Returns the config of the optimizer. Optimizer configuration for AdaBelief. diff --git a/lib/model/optimizers/lion.py b/lib/model/optimizers/lion.py new file mode 100644 index 0000000000..de90119178 --- /dev/null +++ b/lib/model/optimizers/lion.py @@ -0,0 +1,110 @@ +#! /usr/env/bin/python3 +"""PyTorch implementation of the Lion optimizer.""" +# Copyright 2023 Google Research. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +import logging +import typing as T + +import torch +from torch.optim.optimizer import Optimizer + +from lib.logger import parse_class_init +from lib.utils import get_module_objects + +logger = logging.getLogger(__name__) + + +class Lion(Optimizer): + """Lion optimizer from Google + + Parameters + ---------- + params + Iterable of parameters to optimize or dicts defining parameter groups + lr + Learning rate. Default: 1e-4 + betas + Coefficients used for computing running averages of gradient and its square. + Default: (0.9, 0.99) + weight_decay + Weight decay coefficient. Default: 0 + + Reference + --------- + https://github.com/google/automl/blob/master/lion/lion_pytorch.py + """ + def __init__(self, + params: T.Iterable, + lr: float = 1e-4, + betas: tuple[float, float] = (0.9, 0.99), + weight_decay: float = 0.0) -> None: + logger.debug(parse_class_init(locals())) + if 0.0 > lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") + defaults = {"lr": lr, "betas": betas, "weight_decay": weight_decay} + super().__init__(params, defaults) + + @torch.no_grad() + def step(self, closure: T.Callable | None = None) -> torch.Tensor: # type:ignore[override] + """Performs a single optimization step. + + Parameters + ---------- + closure + A closure that reevaluates the model and returns the loss. + + Returns + ------- + The loss + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + + # Perform step weight decay + p.data.mul_(1 - group["lr"] * group["weight_decay"]) + + grad = p.grad + state = self.state[p] + # State initialization + if len(state) == 0: + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like(p) + + exp_avg = state["exp_avg"] + beta1, beta2 = group["betas"] + + # Weight update + update = exp_avg * beta1 + grad * (1 - beta1) + + p.add_(update.sign_(), alpha=-group["lr"]) + + # Decay the momentum running average coefficient + exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2) + + return T.cast(torch.Tensor, loss) + + +__all__ = get_module_objects(__name__) diff --git a/lib/training/__init__.py b/lib/training/__init__.py index a6433c1e1d..e06533d81e 100644 --- a/lib/training/__init__.py +++ b/lib/training/__init__.py @@ -4,8 +4,6 @@ from __future__ import annotations import typing as T -from .lr_finder import LearningRateFinder -from .lr_warmup import LearningRateWarmup from .preview_cv import PreviewBuffer, TriggerType if T.TYPE_CHECKING: diff --git a/lib/training/lr_finder.py b/lib/training/lr_finder.py index c4f528c871..a99b6a251d 100644 --- a/lib/training/lr_finder.py +++ b/lib/training/lr_finder.py @@ -15,11 +15,10 @@ from lib.logger import parse_class_init from lib.utils import get_module_objects -from plugins.train import train_config as cfg if T.TYPE_CHECKING: - import torch - from keras import optimizers + from torch import Tensor + from torch.optim.lr_scheduler import ExponentialLR from . import train logger = logging.getLogger(__name__) @@ -39,40 +38,48 @@ class LearningRateFinder: # pylint:disable=too-many-instance-attributes ---------- trainer The training loop with the loaded training plugin + scheduler + The LRFinder scheduler + steps + The number of steps to run the finder for + strength + How aggressively to set the optimal learning rate + mode + The mode to run the Learning Rate Finder in stop_factor When to stop finding the optimal learning rate beta Amount to smooth loss by, for graphing purposes """ - def __init__(self, # pylint:disable=too-many-positional-arguments + def __init__(self, trainer: train.Trainer, + scheduler: ExponentialLR, + steps: int, + strength: T.Literal["default", "aggressive", "extreme"], + mode: T.Literal["set", "graph_and_set", "graph_and_exit"], stop_factor: int = 4, beta: float = 0.98) -> None: logger.debug(parse_class_init(locals())) - self._iterations = cfg.lr_finder_iterations() - self._save_graph = cfg.lr_finder_mode() in ("graph_and_set", "graph_and_exit") - self._strength = LRStrength[cfg.lr_finder_strength().upper()].value - - self._start_lr = 1e-10 - end_lr = 1e+1 - self._trainer = trainer - - self._model = trainer._plugin.model - self._optimizer = trainer._plugin.model.model.optimizer - + self._scheduler = scheduler + self._steps = steps + self._strength = LRStrength[strength.upper()].value + self._mode = mode self._stop_factor = stop_factor self._beta = beta - self._lr_multiplier: float = (end_lr / self._start_lr) ** (1.0 / self._iterations) - self._metrics: dict[T.Literal["learning_rates", "losses"], list[float]] = { - "learning_rates": [], - "losses": []} + self._model = trainer._plugin.model + self._losses: list[float] = [] + self._learning_rates: list[float] = [] self._loss: dict[T.Literal["avg", "best"], float] = {"avg": 0.0, "best": 1e9} + self._best_lr: None | float = None - logger.debug("Initialized %s", self.__class__.__name__) + @property + def best_lr(self) -> None | float: + """The discovered best learning rate or ``None`` if not found""" + return self._best_lr - def _on_batch_end(self, iteration: int, loss: float) -> None: + def _on_batch_end(self, iteration: int, loss: float) -> bool: """Learning rate actions to perform at the end of a batch Parameters @@ -81,26 +88,29 @@ def _on_batch_end(self, iteration: int, loss: float) -> None: The current iteration loss The loss value for the current batch + + Returns + ------- + ``True`` if training should cease. ``False`` to continue """ - learning_rate = float(self._optimizer.learning_rate.numpy()) - self._metrics["learning_rates"].append(learning_rate) + if np.isnan(loss): + logger.info("Loss has NaN'd. Exiting early") + return True + self._learning_rates.append(T.cast(float, self._scheduler.get_last_lr()[0])) self._loss["avg"] = (self._beta * self._loss["avg"]) + ((1 - self._beta) * loss) smoothed = self._loss["avg"] / (1 - (self._beta ** iteration)) - self._metrics["losses"].append(smoothed) + self._losses.append(smoothed) stop_loss = self._stop_factor * self._loss["best"] - if iteration > 1 and smoothed > stop_loss: - self._model.model.stop_training = True - return + logger.info("Loss has diverged. Exiting early") + return True if iteration == 1 or smoothed < self._loss["best"]: self._loss["best"] = smoothed - learning_rate *= self._lr_multiplier - - self._optimizer.learning_rate.assign(learning_rate) + return False def _update_description(self, progress_bar: tqdm) -> None: """Update the description of the progress bar for the current iteration @@ -110,106 +120,46 @@ def _update_description(self, progress_bar: tqdm) -> None: progress_bar The learning rate finder progress bar to update """ - current = self._metrics['learning_rates'][-1] - best_idx = self._metrics["losses"].index(self._loss["best"]) - best = self._metrics["learning_rates"][best_idx] / self._strength + current = self._learning_rates[-1] + best_idx = self._losses.index(self._loss["best"]) + best = self._learning_rates[best_idx] / self._strength progress_bar.set_description(f"Current: {current:.1e} Best: {best:.1e}") def _train(self) -> None: """Train the model for the given number of iterations to find the optimal learning rate and show progress""" logger.info("Finding optimal learning rate...") - p_bar = tqdm(range(1, self._iterations + 1), + p_bar = tqdm(range(1, self._steps + 1), desc="Current: N/A Best: N/A ", leave=False) for idx in p_bar: loss = self._trainer.train_one_batch() - total_loss = T.cast("torch.Tensor", sum(x.total for x in loss)).item() + total_loss = T.cast("Tensor", sum(x.total for x in loss)).item() - if np.isnan(total_loss): - logger.warning("NaN detected! Exiting early") + if self._on_batch_end(idx, total_loss): + logger.debug("[LearningRateFinder] Exiting early") break - self._on_batch_end(idx, total_loss) - self._update_description(p_bar) - def _rebuild_optimizer(self, optimizer: optimizers.Optimizer) -> optimizers.Optimizer: - """Pass through nested Optimizers (eg LossScaleOptimizer) and create new nested - optimizers based on their original config + self._update_description(p_bar) - Returns - ------- - A new optimizer of the same type as the given one, with the same config - """ - logger.debug("Processing optimizer: '%s'", optimizer.name) - config = optimizer.get_config() - if hasattr(optimizer, "inner_optimizer"): - config["inner_optimizer"] = self._rebuild_optimizer(optimizer.inner_optimizer) - retval = optimizer.__class__(**config) - logger.debug("Created optimizer '%s': (old: %s, new: %s)", - optimizer.name, optimizer, retval) - return retval - - def _reset_model(self, original_lr: float, new_lr: float) -> None: + def _reset_model(self, new_lr: float) -> None: """Reset the model's weights to initial values, reset the model's optimizer and set the learning rate Parameters ---------- - original_lr - The model's original learning rate new_lr The discovered optimal learning rate """ self._model.state.add_lr_finder(new_lr) self._model.state.save() - if cfg.lr_finder_mode() == "graph_and_exit": + if self._mode == "graph_and_exit": return - logger.debug("Resetting optimizer") - optimizer = self._rebuild_optimizer(self._optimizer) - del self._optimizer - del self._model.model.optimizer - logger.info("Loading initial weights") self._model.model.load_weights(self._model.io.filename) - self._model.model.compile(optimizer=optimizer, - loss=self._model.model.loss, - metrics=self._model.model.loss) - - logger.info("Updating Learning Rate from %s to %s", f"{original_lr:.1e}", f"{new_lr:.1e}") - self._model.model.optimizer.learning_rate.assign(new_lr) - self._optimizer = self._model.model.optimizer - - def find(self) -> bool: - """Find the optimal learning rate - - Returns - ------- - ``True`` if the learning rate was successfully discovered otherwise ``False`` - """ - if not self._model.io.model_exists: - self._model.io.save() - - original_lr = float(self._model.model.optimizer.learning_rate.numpy()) - self._model.model.optimizer.learning_rate.assign(self._start_lr) - - self._train() - print("\x1b[2K", end="\r") # Clear line - - best_idx = self._metrics["losses"].index(self._loss["best"]) - new_lr = self._metrics["learning_rates"][best_idx] / self._strength - if new_lr < 1e-9: - logger.error("The optimal learning rate could not be found. This is most likely " - "because you did not run the finder for enough iterations.") - shutil.rmtree(self._model.io.model_dir) - return False - - self._plot_loss() - self._reset_model(original_lr, new_lr) - return True - def _plot_loss(self, skip_begin: int = 10, skip_end: int = 1) -> None: """Plot a graph of loss vs learning rate and save to the training folder @@ -220,15 +170,15 @@ def _plot_loss(self, skip_begin: int = 10, skip_end: int = 1) -> None: skip_end Number of iterations to skip at the end. Default: `1` """ - if not self._save_graph: + if self._mode not in ("graph_and_set", "graph_and_exit"): return matplotlib.use("Agg") - lrs = self._metrics["learning_rates"][skip_begin:-skip_end] - losses = self._metrics["losses"][skip_begin:-skip_end] + lrs = self._learning_rates[skip_begin:-skip_end] + losses = self._losses[skip_begin:-skip_end] plt.plot(lrs, losses, label="Learning Rate") - best_idx = self._metrics["losses"].index(self._loss["best"]) - best_lr = self._metrics["learning_rates"][best_idx] + best_idx = self._losses.index(self._loss["best"]) + best_lr = self._learning_rates[best_idx] for val, color in zip(LRStrength, ("g", "y", "r")): l_r = best_lr / val.value idx = lrs.index(next(r for r in lrs if r >= l_r)) @@ -247,5 +197,25 @@ def _plot_loss(self, skip_begin: int = 10, skip_end: int = 1) -> None: logger.info("Saving Learning Rate Finder graph to: '%s'", output) plt.savefig(output) + def find(self) -> None: + """Find the optimal learning rate""" + if not self._model.io.model_exists: + self._model.io.save() + + self._train() + print("\x1b[2K", end="\r") # Clear line + + best_idx = self._losses.index(self._loss["best"]) + new_lr = self._learning_rates[best_idx] / self._strength + if new_lr < 1e-9: + logger.error("The optimal learning rate could not be found. This is most likely " + "because you did not run the finder for enough iterations.") + shutil.rmtree(self._model.io.model_dir) + return + + self._best_lr = new_lr + self._plot_loss() + self._reset_model(new_lr) + __all__ = get_module_objects(__name__) diff --git a/lib/training/lr_warmup.py b/lib/training/lr_warmup.py index 6bbb33ee7e..9245dc8b0a 100644 --- a/lib/training/lr_warmup.py +++ b/lib/training/lr_warmup.py @@ -1,104 +1,113 @@ #! /usr/env/bin/python3 -""" Handles Learning Rate Warmup when training a model """ +"""Handles Learning Rate Warmup when training a model""" from __future__ import annotations import logging import typing as T +from torch.optim.lr_scheduler import LRScheduler + +from lib.logger import parse_class_init from lib.utils import get_module_objects if T.TYPE_CHECKING: - from keras import models + from torch import Tensor + from torch.optim import Optimizer logger = logging.getLogger(__name__) -class LearningRateWarmup(): - """ Handles the updating of the model's learning rate during Learning Rate Warmup +class WarmupScheduler(LRScheduler): + """Handles the updating of the model's learning rate during Learning Rate Warmup Parameters ---------- - model : :class:`keras.models.Model` - The keras model that is to be trained - target_learning_rate : float - The final learning rate at the end of warmup - steps : int + optimizer + The torch optimizer in use + steps The number of iterations to warmup the learning rate for + last_epoch + The last step that was run (last_epoch is a misnomer inherited from PyTorch and actually + refers to steps in our use case). Default: -1 (not yet started) """ - def __init__(self, model: models.Model, target_learning_rate: float, steps: int) -> None: - self._model = model - self._target_lr = target_learning_rate - self._steps = steps - self._current_lr = 0.0 - self._current_step = 0 - self._reporting_points = [int(self._steps * i / 10) for i in range(11)] - logger.debug("Initialized %s", self) - - def __repr__(self) -> str: - """ Pretty string representation for logging """ - call_args = ", ".join(f"{k}={v}" for k, v in {"model": self._model, - "target_learning_rate": self._target_lr, - "steps": self._steps}.items()) - current_params = ", ".join(f"{k[1:]}: {v}" for k, v in self.__dict__.items() - if k not in ("_model", "_target_lr", "_steps")) - return f"{self.__class__.__name__}({call_args}) [{current_params}]" + def __init__(self, optimizer: Optimizer, steps: int, last_epoch: int = -1) -> None: + logger.debug(parse_class_init(locals())) + self.steps = steps + """The total number of steps to warmup the LR for""" + self._reporting_points = [int(self.steps * i / 10) for i in range(11)] + super().__init__(optimizer, last_epoch) @classmethod - def _format_notation(cls, value: float) -> str: - """ Format a float to scientific notation at 1 decimal place + def _fmt(cls, value: float) -> str: + """Format a float to scientific notation at 1 decimal place Parameters ---------- - value : float + value The value to format Returns ------- - str - The formatted float in scientific notation at 1 decimal place + The formatted float in scientific notation at 1 decimal place """ return f"{value:.1e}" - def _set_learning_rate(self) -> None: - """ Set the learning rate for the current step """ - self._current_lr = self._current_step / self._steps * self._target_lr - self._model.optimizer.learning_rate.assign(self._current_lr) - logger.debug("Learning rate set to %s for step %s/%s", - self._current_lr, self._current_step, self._steps) + def get_lr(self) -> list[float | Tensor]: + """Get the learning rate for the current step + + Returns + ------- + The next learning rate for each parameter group for the next step + """ + if self.last_epoch >= self.steps: + return self.base_lrs + + factor = self.last_epoch / self.steps + lrs = [base_lr * factor for base_lr in self.base_lrs] + logger.trace("Learning rate set to %s for step %s/%s", # type:ignore[attr-defined] + lrs, self.last_epoch, self.steps) + return lrs def _output_status(self) -> None: - """ Output the progress of Learning Rate Warmup at set intervals """ - if self._current_step == 1: + """Output the progress of Learning Rate Warmup at set intervals""" + step = self.last_epoch + if step < 1: + return + + current_lr = T.cast(float, self.get_last_lr()[0]) + target_lr = T.cast(float, self.base_lrs[0]) + + if step == 1: logger.info("[Learning Rate Warmup] Start: %s, Target: %s, Steps: %s", - self._format_notation(self._current_lr), - self._format_notation(self._target_lr), self._steps) + self._fmt(current_lr), self._fmt(target_lr), self.steps) return - if self._current_step == self._steps: + if step == self.steps: print() - logger.info("[Learning Rate Warmup] Final Learning Rate: %s", - self._format_notation(self._target_lr)) + logger.info("[Learning Rate Warmup] Final Learning Rate: %s", self._fmt(target_lr)) return - if self._current_step in self._reporting_points: + if step in self._reporting_points: print() progress = int(round(100 / (len(self._reporting_points) - 1) * - self._reporting_points.index(self._current_step), 0)) + self._reporting_points.index(step), 0)) logger.info("[Learning Rate Warmup] Step: %s/%s (%s), Current: %s, Target: %s", - self._current_step, - self._steps, + step, + self.steps, f"{progress}%", - self._format_notation(self._current_lr), - self._format_notation(self._target_lr)) + self._fmt(current_lr), + self._fmt(target_lr)) - def __call__(self) -> None: - """ If a learning rate update is required, update the model's learning rate, otherwise - do nothing """ - if self._steps == 0 or self._current_step >= self._steps: - return + def step(self, epoch=None) -> None: + """If a learning rate update is required, update the model's learning rate, otherwise + do nothing - self._current_step += 1 - self._set_learning_rate() + Parameters + ---------- + epoch + Deprecated argument from PyTorch that should always be ``None``. Default: ``None`` + """ + super().step(epoch) self._output_status() diff --git a/lib/training/optimizer.py b/lib/training/optimizer.py new file mode 100644 index 0000000000..3fc9d39d35 --- /dev/null +++ b/lib/training/optimizer.py @@ -0,0 +1,497 @@ +#!/usr/bin/env python3 +"""Wraps the selected Torch optimizer and handles optimizer related functions such as loss scaling, +clipping and gradient accumulation""" +from __future__ import annotations + +import logging +import typing as T + +import torch +from torch import nn +from torch.optim.lr_scheduler import ExponentialLR + +from lib.logger import parse_class_init +from lib.model.autoclip import AutoClipper +from lib.model import optimizers +from lib.utils import get_module_objects + +from .lr_finder import LearningRateFinder +from .lr_warmup import WarmupScheduler + +if T.TYPE_CHECKING: + from keras import Model as K_Model, Variable + from plugins.train.model._base import ModelBase as Model + from plugins.train.train_config import Optimizer as OptConfig + from .train import Trainer + + +logger = logging.getLogger(__name__) + +_OPTIMIZERS = {"adabelief": optimizers.AdaBelief, + "adam": torch.optim.Adam, + "adamax": torch.optim.Adamax, + "adamw": torch.optim.AdamW, + "lion": optimizers.Lion, + "nadam": torch.optim.NAdam, + "rms-prop": torch.optim.RMSprop} + + +def get_parameter_group_ids(trainable_variables: list[Variable] + ) -> dict[int, T.Literal["decay", "no_decay"]]: + """Obtain the index of each item in the keras model's trainable weights that belong to each + of the optimizer's parameter groups (ie split by weights that take decay and don't take decay) + + Parameters + ---------- + trainable_variables + list of trainable variables from keras model + + Returns + ------- + dictionary of keras model's trainable weight index to the name of the parameter group + """ + retval: dict[int, T.Literal["decay", "no_decay"]] = {} + for idx, var in enumerate(trainable_variables): + retval[idx] = "no_decay" if var.ndim <= 1 or var.name.endswith("bias") else "decay" + + logger.debug("parameter group ids: %s", retval) + return retval + + +class GradClip: + """Handles the clipping of gradients based on user supplied parameters + + Parameters + ---------- + method + The clipping method to use + value + The clipping value to use. For autoclip this is the percentile to clip at (a value of 1.0 + will clip at the 10th percentile a value of 2.5 will clip at the 25th percentile etc) + autoclip_history + The history length for auto clipping. Default: 10000 + """ + def __init__(self, + method: T.Literal["autoclip", "global_norm", "norm", "value"], + value: float, + autoclip_history: int = 10000) -> None: + logger.debug(parse_class_init(locals())) + self._value = value + self._clipper = self._get_clipper(method, autoclip_history) + + @classmethod + def _clip_norm(cls, parameters: list[nn.Parameter], max_norm: float) -> None: + """Clip each parameter independently by its own norm + + Parameters + ---------- + parameters + The parameters to clip + max_norm + The value to clip by + """ + with torch.no_grad(): + for param in parameters: + if param.grad is None: + continue + grad = param.grad + norm = grad.norm(2) + if norm > max_norm: + grad.mul_(max_norm / norm) + + def _get_clipper(self, + method: T.Literal["autoclip", "global_norm", "norm", "value"], + autoclip_history: int) -> T.Callable[[list[nn.Parameter], float], + None | torch.Tensor]: + """Obtain the correct function to clip the gradients based on the selected method + + Parameters + ---------- + method + The clipping method to use + autoclip_history + The history length for auto clipping + + Returns + ------- + The function used to clip the gradients + """ + methods: dict[str, T.Callable[[list[nn.Parameter], float], None | torch.Tensor]] = { + "autoclip": AutoClipper(int(self._value * 10), history_size=autoclip_history), + "global_norm": nn.utils.clip_grad_norm_, + "norm": self._clip_norm, + "value": nn.utils.clip_grad_value_} + if method not in methods: + raise ValueError(f"'{method}' is not a valid clipping method. Select " + f"from {list(methods)}") + retval = methods[method] + logger.debug("[GradClip] Got clipper '%s': %s", method, retval) + return retval + + def __call__(self, parameters: list[nn.Parameter]) -> None: + """Clip the given parameters by the chosen method + + Parameters + ---------- + parameters + The parameters to clip + """ + self._clipper(parameters, self._value) + + +class Optimizer: + """Object for managing the selected Torch optimizer + + Parameters + ---------- + model + The model that is to be trained + config + The optimizer user configuration options + mixed_precision + ``True`` to train using mixed precision. Default: ``False`` + warmup_steps + The number of steps to warmup the learning rate for. Default: 0 + """ + def __init__(self, + model: Model, + config: type[OptConfig], + mixed_precision: bool = False, + warmup_steps: int = 0) -> None: + logger.debug(parse_class_init(locals())) + self._mixed_precision = mixed_precision + self._accumulation_steps = config.gradient_accumulation() + self._scaler = None if not mixed_precision else torch.amp.grad_scaler.GradScaler() + self._clip = None if config.gradient_clipping() == "none" else GradClip( + T.cast(T.Literal["autoclip", "global_norm", "norm", "value"], + config.gradient_clipping()), + config.clipping_value(), + config.autoclip_history()) + + self._optimizer = self._get_optimizer(model.model, config) + self._warmup = None if warmup_steps < 1 else WarmupScheduler(self._optimizer, warmup_steps) + self._lr_scheduler: ExponentialLR | None = None + + self._load_state(model) + + self._accumulation_count = 0 + self._session_steps = 0 + + @classmethod + def _get_optimizer_kwargs(cls, config: type[OptConfig]) -> dict[str, T.Any]: + """Obtain the keyword arguments for the requested optimizer from the user configuration + + Parameters + ---------- + config + The optimizer user configuration options + + Returns + ------- + The optimizer keyword arguments + """ + retval: dict[str, T.Any] = {"weight_decay": config.weight_decay()} + name = config.optimizer() + + if name != "lion": + retval["eps"] = 10 ** config.epsilon_exponent() + + if name in ("adabelief", "adam", "adamw", "adamax", "lion", "nadam"): + retval["betas"] = (config.ada_beta_1(), config.ada_beta_2()) + + if name in ("adabelief", "adam", "adamw"): + retval["amsgrad"] = config.ada_amsgrad() + + logger.debug("[Optimizer] '%s' kwargs: %s", name, retval) + return retval + + def _get_optimizer(self, model: K_Model, config: type[OptConfig]) -> torch.optim.Optimizer: + """Obtain the configured optimizer the given configuration file options + + Parameters + ---------- + model + The keras model that is to be trained + config + The optimizer user configuration options + + Returns + ------- + The requested configured optimizer + """ + name = config.optimizer() + if name not in _OPTIMIZERS: + raise ValueError(f"'{name}' is not a valid optimizer. Select from {list(_OPTIMIZERS)}") + optimizer = _OPTIMIZERS[name] + + retval = optimizer(self._get_parameter_groups(model, config.weight_decay()), + lr=config.learning_rate(), + **self._get_optimizer_kwargs(config)) + logger.debug("[Optimizer] Got optimizer '%s': %s", name, retval) + return retval + + def _get_parameter_groups(self, model: K_Model, weight_decay: float + ) -> tuple[dict[T.Literal["params", "weight_decay"], + list[nn.Parameter] | float], + dict[T.Literal["params", "weight_decay"], + list[nn.Parameter] | float]]: + """Obtain the parameter groups from within the keras model + + Parameters + ---------- + model + The keras model that is to be trained + weight_decay + The amount of weight decay to apply + + Returns + ------- + The parameters that require weight decay in position 0 and no weight decay in position 1 + """ + index_map = get_parameter_group_ids(model.trainable_variables) + groups: dict[T.Literal["decay", "no_decay"], list[nn.Parameter]] = {"decay": [], + "no_decay": []} + # pylint:disable=protected-access + for idx, var in enumerate(model.trainable_variables): + if not hasattr(var, "_value") or not isinstance(var._value, nn.Parameter): + raise RuntimeError( + f"Cannot extract torch parameter from keras.Variable '{var.name}'. " + "Keras version may have changed internal structure.") + groups[index_map[idx]].append(var._value) + + retval: tuple[dict[T.Literal["params", "weight_decay"], list[nn.Parameter] | float], + dict[T.Literal["params", "weight_decay"], list[nn.Parameter] | float]] = ( + {"params": groups["decay"], "weight_decay": weight_decay}, + {"params": groups["no_decay"], "weight_decay": 0.0} + ) + + logger.debug("[Optimizer] decay params: %s, no_decay params: %s", + {k: len(v) if isinstance(v, list) else v for k, v in retval[0].items()}, + {k: len(v) if isinstance(v, list) else v for k, v in retval[1].items()}) + return retval + + def _from_legacy(self, + state: dict[str, T.Any]) -> dict[str, T.Any] | None: + """Populate the remaining param_group items for weights from legacy saved keras optimizer + and validate shapes + + Parameters + ---------- + state + The partial state_dict migrated from a keras optimizer + + Returns + ------- + The final state_dict grouped for torch or ``None`` if weights could not be mapped + """ + logger.debug("[Optimizer] Loading weights from legacy Keras optimizer") + imported_params = state["optimizer"]["state"] + p_groups = self._optimizer.param_groups + exists = [p for g in p_groups for p in g["params"]] + + if len(imported_params) != len(exists): + logger.warning("Imported optimizer weights count mismatch. Optimizer will be reset") + return None + + for idx, exist in enumerate(exists): + # exp_avg for ada based optimizers, square_avg for rms-prop + key = "exp_avg" if "exp_avg" in imported_params[idx] else "square_avg" + if imported_params[idx][key].shape != exist.shape: + logger.warning("Imported optimizer weights shape mismatch. " + "Optimizer will be reset") + return None + + imported_p_groups = state["optimizer"]["param_groups"] + if len(p_groups) != len(imported_p_groups): + logger.warning("Parameter group count mismatch (exists: %s, imported: %s). " + "Optimizer will be reset", len(p_groups), len(imported_p_groups)) + return None + + for idx, group in enumerate(p_groups): + p_group = state["optimizer"]["param_groups"][idx] + state["optimizer"]["param_groups"][idx] = {k: p_group.get(k, v) + for k, v in group.items()} + + return state + + def load_state_dict(self, state_dict: dict[str, T.Any]) -> None: + """Load the serialized data from a state dict into this object + + Parameters + ---------- + state_dict + The serialized data to load + """ + logger.debug("[Optimizer] Loading state_dict") + self._optimizer.load_state_dict(state_dict["optimizer"]) + if self._scaler is not None and state_dict.get("scaler") is not None: + logger.debug("[Optimizer] Loading scaler state_dict: %s", state_dict["scaler"]) + self._scaler.load_state_dict(state_dict["scaler"]) + + def _load_state(self, model: Model) -> None: + """Load weights if resuming and optimizer weights exist within the model file. + + Also handles migration of legacy Keras optimizer weights to torch optimizer + + Parameters + ---------- + model + The model that is to be trained + """ + if not model.io.model_exists: + logger.debug("[Optimizer] Model file does not exist. Not loading state") + return + + state = model.io.load_optimizer() + if state is None: + logger.debug("[Optimizer] No optimizer saved in model file") + return + + if state["version"] == 0.5: # Migrating from keras optimizer + state = self._from_legacy(state) + if state is None: + return + + self.load_state_dict(state_dict=state) + + def backward(self, loss: torch.Tensor) -> None: + """Perform the optimizer's backward pass + + Parameters + ---------- + loss + The loss scalar from the forward pass + """ + scaled = loss / self._accumulation_steps + if self._scaler: + self._scaler.scale(scaled).backward() + else: + scaled.backward() + + def step(self) -> None: + """Perform the optimizer step if valid and zero the gradients. + + Handles gradient accumulation, scaling for mixed precision and gradient clipping + """ + self._accumulation_count += 1 + if self._accumulation_count != self._accumulation_steps: + return + + if self._clip is not None and self._scaler is not None: + self._scaler.unscale_(self._optimizer) + if self._clip is not None: + self._clip([p for g in self._optimizer.param_groups for p in g["params"]]) + + if self._scaler is None: + self._optimizer.step() + else: + self._scaler.step(self._optimizer) + self._scaler.update() + + if self._lr_scheduler is not None: + self._lr_scheduler.step() + elif self._warmup is not None and self._session_steps < self._warmup.steps: + self._session_steps += 1 + self._warmup.step() + + self._optimizer.zero_grad(set_to_none=True) + self._accumulation_count = 0 + + def state_dict(self) -> dict[str, T.Any]: + """Serialized data as a dict for relevant options contained in this class + + Returns + ------- + The serialized data for this object for saving and loading + """ + return {"version": 1.0, + "optimizer": self._optimizer.state_dict(), + "scaler": None if self._scaler is None else self._scaler.state_dict()} + + def to(self, device: torch.Device) -> None: + """Place the optimizer onto the given device + + Parameters + ---------- + device + The device to place the optimizer on to + """ + logger.debug("[Optimizer] to: %s", device) + for state in self._optimizer.state.values(): + for k, v in state.items(): + if isinstance(v, torch.Tensor): + state[k] = v.to(device) + + def set_lr(self, lr: float) -> None: + """Manually assign the optimizer's learning rate with the given value + + Parameters + ---------- + lr + The learning rate to apply to the optimizer + """ + logger.debug("[Optimizer] Setting learning rate to: %s", lr) + for p in self._optimizer.param_groups: + p["lr"] = lr + if "initial_lr" in p: + p["initial_lr"] = lr + + def find_learning_rate(self, + trainer: Trainer, + steps: int, + start_lr: float, + end_lr: float, + strength: T.Literal["default", "aggressive", "extreme"], + mode: T.Literal["set", "graph_and_set", "graph_and_exit"]) -> bool: + """Use the Learning Rate Finder to discover the optimal learning rate + + Parameters + ---------- + trainer + The training loop with the loaded training plugin + steps + The number of iterations to run the learning rate finder for + start_lr + The learning rate to start scanning from + end_lr + The final learning rate to scan until + strength + How aggressively to set the optimal learning rate + mode + The mode to run the Learning Rate Finder in + + Returns + ------- + ``True`` if an optimal learning rate was discovered. + """ + original_lr = self._optimizer.param_groups[0].get("initial_lr", + self._optimizer.param_groups[0]["lr"]) + self.set_lr(start_lr) + opt_state = self._optimizer.state_dict() + scaler_state = None if self._scaler is None else self._scaler.state_dict() + + gamma: float = (end_lr / start_lr) ** (1.0 / steps) + self._lr_scheduler = ExponentialLR(self._optimizer, gamma=gamma) + + lrf = LearningRateFinder(trainer, self._lr_scheduler, steps, strength, mode) + lrf.find() + + del self._lr_scheduler + self._lr_scheduler = None + + if lrf.best_lr is None: + return False + + logger.debug("[Optimizer] Resetting optimizer for LearningRateFinder: %s", opt_state) + self._optimizer.load_state_dict(opt_state) + if self._scaler is not None and scaler_state is not None: + self._scaler.load_state_dict(scaler_state) + + logger.info("Updating Learning Rate from %s to %s", + f"{original_lr:.1e}", f"{lrf.best_lr:.1e}") + self.set_lr(lrf.best_lr) + + return True + + +__all__ = get_module_objects(__name__) diff --git a/lib/training/train.py b/lib/training/train.py index 6b055e24c4..52658fa6c2 100644 --- a/lib/training/train.py +++ b/lib/training/train.py @@ -1,5 +1,5 @@ #! /usr/env/bin/python3 -"""Run the training loop for a training plugin """ +"""Run the training loop for a training plugin""" from __future__ import annotations import logging @@ -16,7 +16,6 @@ from lib.logger import format_array, parse_class_init from lib.torch_utils import get_device -from lib.training import LearningRateFinder, LearningRateWarmup from lib.training.preview import Samples from lib.training.data import get_label, PreviewLoader, TrainLoader from lib.training.tensorboard import TorchTensorBoard @@ -25,6 +24,7 @@ from plugins.train.trainer import trainer_config as trn_cfg from .loss import LossCollator +from .optimizer import Optimizer if T.TYPE_CHECKING: import numpy.typing as npt @@ -53,6 +53,8 @@ class Trainer: # pylint:disable=too-many-instance-attributes The plugin that will be processing each batch preview ``True`` to generate previews + warmup_steps + The number of steps to warmup the learning rate for. Default: 0 timelapse_folders The input folders to create timelapse images from. Default: ``None`` (no timelapse) timelapse_output @@ -62,6 +64,7 @@ class Trainer: # pylint:disable=too-many-instance-attributes def __init__(self, plugin: TrainerBase, preview: bool, + warmup_steps: int = 0, timelapse_folders: list[str] | None = None, timelapse_output: str = "") -> None: logger.debug(parse_class_init(locals())) @@ -74,19 +77,23 @@ def __init__(self, self._model = plugin.model self._out_size = max(x[1] for x in self._model.output_shapes if x[-1] != 1) self._configure_model(plugin) + self._optimizer = Optimizer(self._model, + mod_cfg.Optimizer, + mixed_precision=mod_cfg.mixed_precision(), + warmup_steps=warmup_steps) + self._optimizer.to(self._device) self._train_loader = self._get_train_loader() - self._preview_loader = self._get_preview_loader() - self._timelapse_loader = self._get_timelapse_loader() self._exit_early = self._handle_lr_finder() if self._exit_early: logger.debug("[Trainer] Exiting from LR Finder") return - self._warmup = self._get_warmup() - self._model.state.add_session_batchsize(plugin.batch_size) + self._preview_loader = self._get_preview_loader() + self._timelapse_loader = self._get_timelapse_loader() + self._model.state.add_session_batchsize(plugin.batch_size) self._tensorboard = self._set_tensorboard() self._samples = Samples(self._model.coverage_ratio, mod_cfg.Loss.learn_mask() or mod_cfg.Loss.penalized_mask_loss(), @@ -227,28 +234,26 @@ def _handle_lr_finder(self) -> bool: learning_rate = self._model.state.lr_finder logger.info("Setting learning rate from Learning Rate Finder to %s", f"{learning_rate:.1e}") - self._model.model.optimizer.learning_rate.assign(learning_rate) + self._optimizer.set_lr(learning_rate) self._model.state.update_session_config("learning_rate", learning_rate) return False if self._model.state.iterations == 0 and self._model.state.session_id == 1: - lrf = LearningRateFinder(self) - success = lrf.find() + success = self._optimizer.find_learning_rate( + self, + mod_cfg.lr_finder_iterations(), + 1e-10, + 1e-1, + T.cast(T.Literal["default", "aggressive", "extreme"], + mod_cfg.lr_finder_strength()), + T.cast(T.Literal["set", "graph_and_set", "graph_and_exit"], + mod_cfg.lr_finder_mode()) + ) return mod_cfg.lr_finder_mode() == "graph_and_exit" or not success logger.debug("[Trainer] No learning rate finder rate. Not setting") return False - def _get_warmup(self) -> LearningRateWarmup: - """Obtain the learning rate warmup instance - - Returns - ------- - The Learning Rate Warmup object - """ - target_lr = float(self._model.model.optimizer.learning_rate.value.cpu().numpy()) - return LearningRateWarmup(self._model.model, target_lr, self._model.warmup_steps) - def _set_tensorboard(self) -> TorchTensorBoard | None: """Set up Tensorboard callback for logging loss. @@ -290,6 +295,7 @@ def train_one_batch(self) -> list[BatchLoss]: inputs, targets, meta = next(self._train_loader) loss = self._plugin.train_batch([i.to(self._device) for i in inputs], [t.to(self._device) for t in targets], + self._optimizer, meta.to(self._device)) retval = [x.to_cpu() for x in loss] except OutOfMemoryError as err: @@ -527,7 +533,6 @@ def train_one_step(self, do_snapshot = (self._plugin.config.snapshot_interval != 0 and self._model.iterations - 1 >= self._plugin.config.snapshot_interval and (self._model.iterations - 1) % self._plugin.config.snapshot_interval == 0) - self._warmup() loss = self.train_one_batch() self._log_tensorboard(loss) total_loss = self._collate_and_store_loss(loss) @@ -555,7 +560,7 @@ def save(self, is_exit: bool = False) -> None: is_exit ``True`` if save has been called on model exit. Default: ``False`` """ - self._model.io.save(is_exit=is_exit) + self._model.io.save(self._optimizer, is_exit=is_exit) assert self._tensorboard is not None self._tensorboard.on_save() if is_exit: diff --git a/plugins/extract/detect/mtcnn.py b/plugins/extract/detect/mtcnn.py index d4440b4ceb..4b226e40be 100644 --- a/plugins/extract/detect/mtcnn.py +++ b/plugins/extract/detect/mtcnn.py @@ -170,7 +170,7 @@ class PNet(nn.Module): Parameters ---------- weights_path - The path to the keras model file + The path to the torch model file """ def __init__(self, weights_path: str) -> None: super().__init__() @@ -217,7 +217,7 @@ class PNetRunner(): Parameters ---------- weights_path - The path to the keras model file + The path to the torch model file device The device to use for model inference input_size @@ -419,7 +419,7 @@ class RNetRunner(): Parameters ---------- weights_path - The path to the keras model file + The path to the torch model file device The device to run inference on input_size @@ -575,7 +575,7 @@ class ONetRunner(): Parameters ---------- weights_path - The path to the keras model file + The path to the torch model file device The device to run inference on input_size diff --git a/plugins/train/model/_base/io.py b/plugins/train/model/_base/io.py index a4635f2a7d..518d25036d 100644 --- a/plugins/train/model/_base/io.py +++ b/plugins/train/model/_base/io.py @@ -1,6 +1,5 @@ #!/usr/bin/env python3 -""" -IO handling for the model base plugin. +"""IO handling for the model base plugin. The objects in this module should not be called directly, but are called from :class:`~plugins.train.model._base.ModelBase` @@ -10,69 +9,77 @@ - The loading and freezing of weights for model plugins. """ from __future__ import annotations +import gc +import io +import json import logging import os import sys import typing as T -from keras import layers, models as kmodels +import zipfile + +from keras import layers, models as k_models, Variable +import numpy as np +import torch from lib.logger import parse_class_init from lib.model.backup_restore import Backup +from lib.training.optimizer import get_parameter_group_ids from lib.utils import get_module_objects, FaceswapError from .update import Legacy, PatchKerasConfig if T.TYPE_CHECKING: + from keras.optimizers import Optimizer as K_Optimizer, LossScaleOptimizer + from lib.training.optimizer import Optimizer from .model import ModelBase - from keras import Optimizer logger = logging.getLogger(__name__) def get_all_sub_models( - model: kmodels.Model, - models: list[kmodels.Model] | None = None) -> list[kmodels.Model]: - """ For a given model, return all sub-models that occur (recursively) as children. + model: k_models.Model, + models: list[k_models.Model] | None = None) -> list[k_models.Model]: + """For a given model, return all sub-models that occur (recursively) as children. Parameters ---------- - model: :class:`keras.models.Model` + model A Keras model to scan for sub models - models: `None` + models Do not provide this parameter. It is used for recursion Returns ------- - list - A list of all :class:`keras.models.Model` objects found within the given model. - The provided model will always be returned in the first position + A list of all :class:`keras.models.Model` objects found within the given model. The provided + model will always be returned in the first position """ if models is None: models = [model] else: models.append(model) for layer in model.layers: - if isinstance(layer, kmodels.Model): + if isinstance(layer, k_models.Model): get_all_sub_models(layer, models=models) return models class IO(): - """ Model saving and loading functions. + """Model saving and loading functions. Handles the loading and saving of the plugin model from disk as well as the model backup and snapshot functions. Parameters ---------- - plugin: :class:`Model` + plugin The parent plugin class that owns the IO functions. - model_dir: str + model_dir The full path to the model save location - is_predict: bool + is_predict ``True`` if the model is being loaded for inference. ``False`` if the model is being loaded for training. - save_optimizer: ["never", "always", "exit"] + save_optimizer When to save the optimizer weights. `"never"` never saves the optimizer weights. `"always"` always saves the optimizer weights. `"exit"` only saves the optimizer weights on an exit request. @@ -86,69 +93,65 @@ def __init__(self, self._plugin = plugin self._is_predict = is_predict self._model_dir = model_dir - self._save_optimizer = save_optimizer + self._do_save_optimizer = save_optimizer self._history: list[float] = [] - """list[float]: Loss history for current save iteration """ + """Loss history for current save iteration""" self._backup = Backup(self._model_dir, self._plugin.name) self._update_legacy() - logger.debug("Initialized %s", self.__class__.__name__) @property def model_dir(self) -> str: - """ str: The full path to the model folder """ + """The full path to the model folder""" return self._model_dir @property def filename(self) -> str: - """str: The filename for this model.""" + """The filename for this model.""" return os.path.join(self._model_dir, f"{self._plugin.name}.keras") @property def model_exists(self) -> bool: - """ bool: ``True`` if a model of the type being loaded exists within the model folder - location otherwise ``False``. - """ + """``True`` if a model of the type being loaded exists within the model folder location + otherwise ``False``.""" return os.path.isfile(self.filename) @property def history(self) -> list[float]: - """ list[float]: list of loss history for the current save iteration. """ + """list of loss history for the current save iteration.""" return self._history @property def multiple_models_in_folder(self) -> list[str] | None: - """ :list: or ``None`` If there are multiple model types in the requested folder, or model - types that don't correspond to the requested plugin type, then returns the list of plugin - names that exist in the folder, otherwise returns ``None`` """ + """If there are multiple model types in the requested folder, or model types that don't + correspond to the requested plugin type, then returns the list of plugin names that exist + in the folder, otherwise returns ``None``""" plugins = [fname.replace(".keras", "") for fname in os.listdir(self._model_dir) if fname.endswith(".keras")] test_names = plugins + [self._plugin.name] test = False if not test_names else os.path.commonprefix(test_names) == "" retval = None if not test else plugins - logger.debug("plugin name: %s, plugins: %s, test result: %s, retval: %s", + logger.debug("[IO] plugin name: %s, plugins: %s, test result: %s, retval: %s", self._plugin.name, plugins, test, retval) return retval def _update_legacy(self) -> None: - """ Look for faceswap 2.x .h5 files in the model folder. If exists, then update to Faceswap - 3 .keras file and backup the original model .h5 file - - Note: Currently disabled as keras hangs trying to load old faceswap models - """ + """Look for faceswap 2.x .h5 files in the model folder. If exists, then update to Faceswap + 3 .keras file and backup the original model .h5 file""" if self.model_exists: - logger.debug("Existing model file is current: '%s'", os.path.basename(self.filename)) + logger.debug("[IO] Existing model file is current: '%s'", + os.path.basename(self.filename)) return old_fname = f"{os.path.splitext(self.filename)[0]}.h5" if not os.path.isfile(old_fname): - logger.debug("No legacy model file to update") + logger.debug("[IO] No legacy model file to update") return Legacy(old_fname) - def load(self) -> kmodels.Model: - """ Loads the model from disk + def load(self) -> k_models.Model: + """Loads the model from disk If the predict function is to be called and the model cannot be found in the model folder then an error is logged and the process exits. @@ -158,16 +161,15 @@ def load(self) -> kmodels.Model: Returns ------- - :class:`keras.models.Model` - The saved model loaded from disk + The saved model loaded from disk """ - logger.debug("Loading model: %s", self.filename) + logger.debug("[IO] Loading model: %s", self.filename) if self._is_predict and not self.model_exists: logger.error("Model could not be found in folder '%s'. Exiting", self._model_dir) sys.exit(1) try: - model = kmodels.load_model(self.filename, compile=False) + model = k_models.load_model(self.filename, compile=False) except RuntimeError as err: if "unable to get link info" in str(err).lower(): msg = (f"Unable to load the model from '{self.filename}'. This may be a " @@ -199,69 +201,89 @@ def load(self) -> kmodels.Model: logger.info("Loaded model from disk: '%s'", self.filename) return model # pyright:ignore[reportReturnType] - def _remove_optimizer(self) -> Optimizer: - """ Keras 3 `.keras` format ignores the `save_optimizer` kwarg. To hack around this we - remove the optimizer from the model prior to saving and then re-attach it to the model + def load_optimizer(self) -> dict[str, T.Any] | None: + """Load the optimizer's state_dict from the .keras model file Returns ------- - :class:`keras.optimizers.Optimizer` | None - The optimizer for the model, if it should not be saved. ``None`` if it should be saved + The saved optimizer state_dict or ``None`` if it does not exist """ - retval = self._plugin.model.optimizer - del self._plugin.model.optimizer - logger.debug("Removed optimizer for saving: %s", retval) + logger.debug("[IO] Loading optimizer state_dict") + opt_file = "optimizer.pt" + keras_conf = "config.json" + with zipfile.ZipFile(self.filename, "r") as z_file: + f_list = z_file.namelist() + if opt_file in f_list: # Saved torch optimizer + retval = torch.load(io.BytesIO(z_file.read(opt_file))) + elif keras_conf in f_list: # convert legacy keras optimizer + conf = json.loads(z_file.read(keras_conf)) + retval = OptimizerMigrate(conf, self.filename).convert() + else: + retval = None + + if retval is None: + logger.debug("[IO] No optimizer in .keras file") + return None + + logger.debug("[IO] Loaded optimizer state_dict: %s", + {k: list(v) if isinstance(v, dict) else v for k, v in retval.items()}) return retval - def _save_model(self, is_exit: bool, force_save_optimizer: bool) -> None: - """ Save the model either with or without the optimizer weights + def _save_optimizer(self, optimizer: Optimizer) -> None: + """Inject the optimizer's state_dict into the .keras model file - Keras 3 ignores 'save_optimizer` so if it should not be saved, we remove it from - the model for saving, then re-attach it + Parameters + ---------- + optimizer + The current optimizer in use for the model that is to be injected + """ + logger.debug("[IO] Saving optimizer: %s", optimizer) + buf = io.BytesIO() + torch.save(optimizer.state_dict(), buf) + opt_bytes = buf.getvalue() + with zipfile.ZipFile(self.filename, "a") as z_file: + z_file.writestr("optimizer.pt", + opt_bytes, + compress_type=zipfile.ZIP_DEFLATED, + compresslevel=1) + + def _save_model(self, optimizer: Optimizer | None, is_exit: bool) -> None: + """Save the model either with or without the optimizer weights Parameters ---------- - is_exit: bool + optimizer + The current optimizer in use for the model if it should be saved + is_exit ``True`` if the save request has come from an exit process request otherwise ``False``. - force_save_optimizer: bool - ``True`` to force saving the optimizer weights with the model, otherwise ``False``. """ - include_optimizer = (force_save_optimizer or - self._save_optimizer == "always" or - (self._save_optimizer == "exit" and is_exit)) - - optimizer = None - if not include_optimizer: - optimizer = self._remove_optimizer() + include_optimizer = (self._do_save_optimizer == "always" or + (self._do_save_optimizer == "exit" and is_exit)) self._plugin.model.save(self.filename) + if include_optimizer and optimizer is not None: + self._save_optimizer(optimizer) self._plugin.state.save() - if not include_optimizer: - assert optimizer is not None - logger.debug("Re-attaching optimizer: %s", optimizer) - setattr(self._plugin.model, "optimizer", optimizer) - def _get_save_average(self) -> float: - """ Return the average loss since the last save iteration and reset historical loss + """Return the average loss since the last save iteration and reset historical loss Returns ------- - float - The average loss since the last save iteration + The average loss since the last save iteration """ - logger.debug("Getting save averages") + logger.debug("[IO] Getting save averages") if not self._history: - logger.debug("No loss in history") + logger.debug("[IO] No loss in history") retval = 0.0 else: retval = sum(self._history) / len(self._history) self._history = [] # Reset historical loss - logger.debug("Average loss since last save: %s", round(retval, 5)) + logger.debug("[IO] Average loss since last save: %s", round(retval, 5)) return retval def _should_backup(self, save_average: float) -> bool: - """ Check whether the loss average for this save iteration is the lowest that has been + """Check whether the loss average for this save iteration is the lowest that has been seen. This protects against model corruption by only backing up the model if the sum of all loss @@ -272,15 +294,15 @@ def _should_backup(self, save_average: float) -> bool: This is by no means a perfect system. If the model corrupts at an iteration close to a save iteration, then the averages may still be pushed lower than a previous save average, resulting in backing up a corrupted model. Changing loss weighting can also - arteficially impact this + artificially impact this Parameters ---------- - save_average: float + save_average The average loss since the last save iteration """ if not self._plugin.state.lowest_avg_loss: - logger.debug("Set initial save iteration loss average: %s", save_average) + logger.debug("[IO] Set initial save iteration loss average: %s", save_average) self._plugin.state.lowest_avg_loss = save_average return False @@ -289,56 +311,53 @@ def _should_backup(self, save_average: float) -> bool: if backup: # Update lowest loss values to the state file self._plugin.state.lowest_avg_loss = save_average - logger.debug("Updated lowest historical save iteration average from: %s to: %s", + logger.debug("[IO] Updated lowest historical save iteration average from: %s to: %s", old_average, save_average) - logger.debug("Should backup: %s", backup) + logger.debug("[IO] Should backup: %s", backup) return backup def _maybe_backup(self) -> tuple[float, bool]: - """ Backup the model if total average loss has dropped for the save iteration + """Backup the model if total average loss has dropped for the save iteration Returns ------- - float + average_loss The total loss average since the last save iteration - bool + backed_up ``True`` if the model was backed up """ save_average = self._get_save_average() should_backup = self._should_backup(save_average) if not save_average or not should_backup: - logger.debug("Not backing up model (save_average: %s, should_backup: %s)", + logger.debug("[IO] Not backing up model (save_average: %s, should_backup: %s)", save_average, should_backup) return save_average, False - logger.debug("Backing up model") + logger.debug("[IO] Backing up model") self._backup.backup_model(self.filename) self._backup.backup_model(self._plugin.state.filename) return save_average, True - def save(self, - is_exit: bool = False, - force_save_optimizer: bool = False) -> None: - """ Backup and save the model and state file. + def save(self, optimizer: Optimizer | None = None, is_exit: bool = False) -> None: + """Backup and save the model and state file. Parameters ---------- - is_exit: bool, optional + optimizer + The current optimizer in use for the model if it should be saved. Default: ``None`` + is_exit ``True`` if the save request has come from an exit process request otherwise ``False``. Default: ``False`` - force_save_optimizer: bool, optional - ``True`` to force saving the optimizer weights with the model, otherwise ``False``. - Default:``False`` """ - logger.debug("Backing up and saving models") + logger.debug("[IO] Backing up and saving models") print("\x1b[2K", end="\r") # Clear last line logger.info("Saving Model...") - self._save_model(is_exit, force_save_optimizer) + self._save_model(optimizer, is_exit) save_average, backed_up = self._maybe_backup() - msg = "[Saved optimizer state for Snapshot]" if force_save_optimizer else "[Saved model]" + msg = "[Saved model]" if save_average: msg += f" - Average total loss since last save: {save_average:.5f}" if backed_up: @@ -346,28 +365,28 @@ def save(self, logger.info(msg) def snapshot(self) -> None: - """ Perform a model snapshot. + """Perform a model snapshot. Notes ----- Snapshot function is called 1 iteration after the model was saved, so that it is built from the latest save, hence iteration being reduced by 1. """ - logger.debug("Performing snapshot. Iterations: %s", self._plugin.iterations) + logger.debug("[IO] Performing snapshot. Iterations: %s", self._plugin.iterations) self._backup.snapshot_models(self._plugin.iterations - 1) - logger.debug("Performed snapshot") + logger.debug("[IO] Performed snapshot") class Weights(): - """ Handling of freezing and loading model weights + """Handling of freezing and loading model weights Parameters ---------- - plugin: :class:`Model` + plugin The parent plugin class that owns the IO functions. """ def __init__(self, plugin: ModelBase) -> None: - logger.debug("Initializing %s: (plugin: %s)", self.__class__.__name__, plugin) + logger.debug(parse_class_init(locals())) self._model = plugin.model self._name = plugin.model_name self._do_freeze = plugin._args.freeze_weights @@ -375,24 +394,22 @@ def __init__(self, plugin: ModelBase) -> None: self._freeze_layers = plugin.freeze_layers self._load_layers = plugin.load_layers - logger.debug("Initialized %s", self.__class__.__name__) @classmethod def _check_weights_file(cls, weights_file: str) -> str | None: - """ Validate that we have a valid path to a .keras file. + """Validate that we have a valid path to a .keras file. Parameters ---------- - weights_file: str + weights_file The full path to a weights file Returns ------- - str - The full path to a weights file + The full path to a weights file """ if not weights_file: - logger.debug("No weights file selected.") + logger.debug("[Weights] No weights file selected.") return None msg = "" @@ -410,7 +427,7 @@ def _check_weights_file(cls, weights_file: str) -> str | None: return weights_file def freeze(self) -> None: - """ If freeze has been selected in the cli arguments, then freeze those models indicated + """If freeze has been selected in the cli arguments, then freeze those models indicated in the plugin's configuration. """ # Blanket unfreeze layers, as checking the value of :attr:`layer.trainable` appears to # return ``True`` even when the weights have been frozen @@ -418,7 +435,7 @@ def freeze(self) -> None: layer.trainable = True if not self._do_freeze: - logger.debug("Freeze weights deselected. Not freezing") + logger.debug("[Weights] Freeze weights deselected. Not freezing") return for layer in get_all_sub_models(self._model): @@ -431,15 +448,15 @@ def freeze(self) -> None: "model: %s", self._freeze_layers) def load(self, model_exists: bool) -> None: - """ Load weights for newly created models, or output warning for pre-existing models. + """Load weights for newly created models, or output warning for pre-existing models. Parameters ---------- - model_exists: bool + model_exists ``True`` if a model pre-exists and is being resumed, ``False`` if this is a new model """ if not self._weights_file: - logger.debug("No weights file provided. Not loading weights.") + logger.debug("[Weights] No weights file provided. Not loading weights.") return if model_exists and self._weights_file: logger.warning("Ignoring weights file '%s' as this model is resuming.", @@ -474,7 +491,7 @@ def load(self, model_exists: bool) -> None: del weights_models if loaded_ops == 0: - raise FaceswapError(f"No weights were succesfully loaded from your weights file: " + raise FaceswapError(f"No weights were successfully loaded from your weights file: " f"'{self._weights_file}'. Please check and try again.") if skipped_ops > 0: logger.warning("%s weight(s) were unable to be loaded for your model. This is most " @@ -482,13 +499,12 @@ def load(self, model_exists: bool) -> None: "different settings than you have set for your current model.", skipped_ops) - def _get_weights_model(self) -> list[kmodels.Model]: - """ Obtain a list of all sub-models contained within the weights model. + def _get_weights_model(self) -> list[k_models.Model]: + """Obtain a list of all sub-models contained within the weights model. Returns ------- - list - List of all models contained within the .keras file + List of all models contained within the .keras file Raises ------ @@ -496,7 +512,7 @@ def _get_weights_model(self) -> list[kmodels.Model]: In the event of a failure to load the weights, or the weights belonging to a different model """ - retval = get_all_sub_models(kmodels.load_model( # pyright:ignore[reportArgumentType] + retval = get_all_sub_models(k_models.load_model( # pyright:ignore[reportArgumentType] self._weights_file, compile=False)) if not retval: @@ -511,26 +527,25 @@ def _load_layer_weights(self, layer: layers.Layer, sub_weights: layers.Layer, model_name: str) -> T.Literal[-1, 0, 1]: - """ Load the weights for a single layer. + """Load the weights for a single layer. Parameters ---------- - layer: :class:`keras.layers.Layer` + layer The layer to set the weights for - sub_weights: list + sub_weights The list of layers in the weights model to load weights from - model_name: str + model_name The name of the current sub-model that is having it's weights loaded Returns ------- - int - `-1` if the layer has no weights to load. `0` if weights loading was unsuccessful. `1` - if weights loading was successful + `-1` if the layer has no weights to load. `0` if weights loading was unsuccessful. `1` if + weights loading was successful """ old_weights = layer.get_weights() if not old_weights: - logger.debug("Skipping layer without weights: %s", layer.name) + logger.debug("[Weights] Skipping layer without weights: %s", layer.name) return -1 layer_weights = next((lyr for lyr in sub_weights.layers @@ -550,4 +565,205 @@ def _load_layer_weights(self, return 1 +class OptimizerMigrate: + """Migrates weights from a keras optimizer to a torch optimizer's state dict""" + def __init__(self, config: dict[str, T.Any], model_path: str): + logger.debug(parse_class_init(locals())) + self._config = config + self._model_path = model_path + ada_map = (("_momentums", "_velocities"), ("exp_avg", "exp_avg_sq")) + self._mapping: dict[str, tuple[tuple[str, ...], tuple[str, ...]]] = { + "AdaBeliefOptimizer": (ada_map[0], ("exp_avg", "exp_avg_var")), + "adam": ada_map, + "adamax": (("_m", "_u"), ("exp_avg", "exp_inf")), + "adamw": ada_map, + "lion": (("_momentums", ), ("exp_avg", )), + "nadam": (ada_map[0] + ("_u_product", ), ada_map[1] + ("mu_product", )), + "rmsprop": (("_velocities", ), ("square_avg",)) + } + + def _get_optimizer_and_group_ids(self) -> tuple[K_Optimizer, + dict[int, + T.Literal["decay", "no_decay"]]] | None: + """Obtain the optimizer from the saved .keras model + + Returns + ------- + optimizer + The saved keras optimizer if it exists or ``None`` if it does not + group_ids + dictionary of keras model's trainable weight index to the name of the parameter group + """ + compile_conf = self._config.get("compile_config", {}).get("optimizer") + if not compile_conf: + logger.debug("[OptimizerMigrate] No saved keras optimizer in model file") + return None + tmp_model = T.cast(k_models.Model, k_models.load_model(self._model_path, compile=True)) + opt = T.cast("K_Optimizer", tmp_model.optimizer) + group_ids = get_parameter_group_ids(tmp_model.trainable_variables) + del tmp_model + gc.collect() + logger.debug("[OptimizerMigrate] keras optimizer from model file: %s", opt) + return opt, group_ids + + def _build_optimizer_state(self, + optimizer: K_Optimizer, + decay_indices: list[int], + no_decay_indices: list[int]) -> dict[int, dict[str, torch.Tensor]]: + """Build the "state" item for the optimizer state_dict + + Parameters + ---------- + optimizer + The loaded keras optimizer + decay_indices + The list of keras variable indices that belong to the decay parameter group + no_decay_indices + The list of keras variable indices that belong to the no_decay parameter group + + Returns + ------- + The populated, ordered, state item in torch format from the keras optimizer + """ + mapping = self._mapping[optimizer.name] + logger.debug("[OptimizerMigrate] mapping for '%s': %s -> %s", + optimizer.name, mapping[0], mapping[1]) + if not all(hasattr(optimizer, x) for x in mapping[0]): + raise RuntimeError( + f"Cannot extract {mapping[0]} from keras optimizer. Keras version may have " + "changed internal structure.") + + if optimizer.name == "lion": + step = {} + else: + step = {"step": torch.from_numpy( + T.cast(np.ndarray, optimizer.iterations.numpy()).astype(np.float32))} + ordered = decay_indices + no_decay_indices + + # pylint:disable=protected-access + vars_ = {mapping[1][idx]: getattr(optimizer, x)._value.data + for idx, x in enumerate(mapping[0]) + if isinstance(getattr(optimizer, x), Variable)} + weights = {x: getattr(optimizer, mapping[0][idx]) + for idx, x in enumerate(mapping[1]) + if x not in vars_} + + retval: dict[int, dict[str, torch.Tensor]] = {} + for dst_idx, src_idx in enumerate(ordered): + layer = {k: v[src_idx] for k, v in weights.items()} + if not all(hasattr(v, "_value") for v in layer.values()): + logger.debug("[OptimizerMigrate] Skipping variable without torch param: %s", + list(layer.values())[0].name.rsplit("_", maxsplit=1)[0]) + continue + c_step = {k: v.clone() for k, v in step.items()} + c_vars = {k: v.clone() for k, v in vars_.items()} + retval[dst_idx] = c_step | c_vars | {k: v._value.data for k, v in layer.items()} + + return retval + + @classmethod + def _get_parameter_groups(cls, + optimizer: K_Optimizer, + weight_indices: list[int], + bias_indices: list[int]) -> list[dict[str, T.Any]]: + """Obtain the fixed config optimizer value and param ids for each parameter group + + Parameters + ---------- + optimizer + The loaded keras optimizer + weight_indices + The list of keras variable indices that belong to the weight parameter group + bias_indices + The list of keras variable indices that belong to the bias parameter group + + Returns + ------- + The parameter group fixed config items and parameter ids + """ + fixed = {} + if hasattr(optimizer, "beta_1") and hasattr(optimizer, "beta_2"): + fixed["betas"] = (optimizer.beta_1, optimizer.beta_2) + if hasattr(optimizer, "amsgrad"): + fixed["amsgrad"] = optimizer.amsgrad + + g1_len = len(weight_indices) + params = [{"params": list(range(g1_len))}, + {"params": list(range(g1_len, g1_len + len(bias_indices)))}] + + retval = [fixed | params[0], fixed | params[1]] + logger.debug("[OptimizerMigrate] param_groups: %s", retval) + return retval + + @classmethod + def _get_scaler_state(cls, + optimizer: LossScaleOptimizer | None) -> dict[str, float | int] | None: + """Build the scaler state_dict from Keras' LossScaleOptimizer + + Parameters + ---------- + optimizer + The Keras LossScaleOptimizer or ``None`` if the optimizer is not scaled + + Returns + ------- + The state dict for Torch scaler or ``None`` if the optimizer is not scaled + """ + if optimizer is None: + logger.debug("[OptimizerMigrate] No scaler to migrate") + return None + + if (not hasattr(optimizer, "dynamic_growth_steps") + or not hasattr(optimizer, "dynamic_scale") + or not hasattr(optimizer, "step_counter")): + logger.warning("Unable to migrate Loss Scaler parameters. Scaler will be reset") + return None + + retval = {"scale": float(optimizer.dynamic_scale.numpy()), + "growth_factor": 2.0, + "backoff_factor": 0.5, + "growth_interval": optimizer.dynamic_growth_steps, + "_growth_tracker": int(optimizer.step_counter.numpy())} + logger.debug("[OptimizerMigrate] scaler: %s", retval) + return retval + + def convert(self) -> dict[str, T.Any] | None: + """Convert the keras optimizer from a keras model file into a torch optimizer state dict + + Returns + ------- + The optimizer state dict for loading into a torch optimizer or ``None`` if no saved + optimizer exists + """ + optimizer_group_ids = self._get_optimizer_and_group_ids() + if optimizer_group_ids is None: + return None + optimizer, index_map = optimizer_group_ids + + scaler_opt: LossScaleOptimizer | None = None + if hasattr(optimizer, "inner_optimizer"): + logger.debug("[OptimizerMigrate] Extracting inner optimizer %s from %s", + optimizer.inner_optimizer, optimizer) + scaler_opt = optimizer + optimizer = optimizer.inner_optimizer + + logger.info("Migrating optimizer weights to Torch") + + weight_indices = [k for k, v in index_map.items() if v == "decay"] + bias_indices = [k for k, v in index_map.items() if v == "no_decay"] + + opt_state = self._build_optimizer_state(optimizer, weight_indices, bias_indices) + if not opt_state: + logger.warning("Unable to migrate optimizer weights. Optimizer will be reset") + return None + + param_groups = self._get_parameter_groups(optimizer, weight_indices, bias_indices) + scaler_state = self._get_scaler_state(scaler_opt) + + retval = {"version": 0.5, + "optimizer": {"state": opt_state, "param_groups": param_groups}, + "scaler": scaler_state} + return retval + + __all__ = get_module_objects(__name__) diff --git a/plugins/train/model/_base/model.py b/plugins/train/model/_base/model.py index c747505e9c..9433116d3d 100644 --- a/plugins/train/model/_base/model.py +++ b/plugins/train/model/_base/model.py @@ -17,7 +17,7 @@ from .inference import Inference from .io import IO, get_all_sub_models, Weights -from .settings import Optimizer, Settings +from .settings import Settings from .state import State if T.TYPE_CHECKING: @@ -139,11 +139,6 @@ def iterations(self) -> int: """The total number of iterations that the model has trained.""" return self._state.iterations - @property - def warmup_steps(self) -> int: - """The number of steps to perform learning rate warmup""" - return self._args.warmup - @property def freeze_layers(self) -> list[str]: """Override to set plugin specific layers that can be frozen. Defaults to ["encoder"]""" @@ -294,20 +289,15 @@ def _output_summary(self) -> None: parent.summary(print_fn=print_fn) def _compile_model(self) -> None: - """Compile the model to include the Optimizer and Loss Function(s).""" + """Legacy from Keras code. Now just load and freeze weights""" logger.debug("Compiling Model") if self.state.model_needs_rebuild: self._model = self._settings.check_model_precision(self._model, self._state) - optimizer = Optimizer().optimizer - if self._settings.use_mixed_precision: - optimizer = self._settings.loss_scale_optimizer(optimizer) - weights = Weights(self) weights.load(self._io.model_exists) weights.freeze() - self.model.compile(optimizer=optimizer) logger.debug("Compiled Model: %s", self.model) def add_history(self, loss: np.ndarray) -> None: diff --git a/plugins/train/model/_base/settings.py b/plugins/train/model/_base/settings.py index 5887ecb65f..ec64fff067 100644 --- a/plugins/train/model/_base/settings.py +++ b/plugins/train/model/_base/settings.py @@ -16,12 +16,8 @@ import keras from keras import config as k_config, dtype_policies, optimizers -from lib.model.optimizers import AdaBelief -from lib.model.autoclip import AutoClipper from lib.model.nn_blocks import reset_naming -from lib.logger import parse_class_init from lib.utils import get_module_objects -from plugins.train.train_config import Optimizer as cfg_opt if T.TYPE_CHECKING: from collections.abc import Callable @@ -31,153 +27,6 @@ logger = logging.getLogger(__name__) -class Optimizer(): - """Obtain the selected optimizer with the appropriate keyword arguments.""" - def __init__(self) -> None: - logger.debug(parse_class_init(locals())) - betas = {"ada_beta_1": "beta_1", "ada_beta_2": "beta_2"} - amsgrad = {"ada_amsgrad": "amsgrad"} - self._valid: dict[str, tuple[T.Type[Optimizer], dict[str, T.Any]]] = { - "adabelief": (AdaBelief, betas | amsgrad), - "adam": (optimizers.Adam, betas | amsgrad), - "adamax": (optimizers.Adamax, betas), - "adamw": (optimizers.AdamW, betas | amsgrad), - "lion": (optimizers.Lion, betas), - "nadam": (optimizers.Nadam, betas), - "rms-prop": (optimizers.RMSprop, {})} - - self._optimizer = self._valid[cfg_opt.optimizer()][0] - self._kwargs: dict[str, T.Any] = {"learning_rate": cfg_opt.learning_rate()} - if cfg_opt.optimizer() != "lion": - self._kwargs["epsilon"] = 10 ** int(cfg_opt.epsilon_exponent()) - - self._configure() - logger.info("Using %s optimizer", self._optimizer.__name__) - logger.debug("Initialized: %s", self.__class__.__name__) - - @property - def optimizer(self) -> optimizers.Optimizer: - """The requested optimizer.""" - return T.cast(optimizers.Optimizer, self._optimizer(**self._kwargs)) - - def _configure_clipping(self, - method: T.Literal["autoclip", "norm", "value", "none"], - value: float, - history: int) -> None: - """Configure optimizer clipping related kwargs, if selected - - Parameters - ---------- - method - The clipping method to use. ``None`` for no clipping - value - The value to clip by norm/value by. For autoclip, this is the clip percentile - (a value of 1.0 is a clip percentile of 10%) - history - autoclip only: The number of iterations to keep for calculating the normalized value - """ - logger.debug("method: '%s', value: %s, history: %s", method, value, history) - if method == "none": - logger.debug("clipping disabled") - return - - logger.info("Enabling Clipping: %s", method.replace("_", " ").replace("_", " ").title()) - clip_types = {"global_norm": "global_clipnorm", "norm": "clipnorm", "value": "clipvalue"} - if method in clip_types: - self._kwargs[clip_types[method]] = value - logger.debug("Setting clipping kwargs for '%s': %s", - method, {k: v for k, v in self._kwargs.items() - if k == clip_types[method]}) - return - - assert method == "autoclip" - # Test for if keras optimizer changes its structure to no longer have _clip_gradients. - # Ensures any tests fails in this situation - assert hasattr(self._optimizer, - "_clip_gradients"), "keras.BaseOptimizer._clip_gradients no longer exists" - - # TODO Keras3 has removed the ""gradient_transformers" kwarg, and there now appears to be - # no standardized method to add custom gradient transformers. Currently, we monkey patch - # its _clip_gradients function, which feels hacky and potentially problematic - setattr(self._optimizer, "_clip_gradients", AutoClipper(int(value * 10), - history_size=history)) - - def _configure_ema(self, enable: bool, momentum: float, frequency: int) -> None: - """configure the optimizer kwargs for exponential moving average updates - - Parameters - ---------- - enable - ``False`` to disable - momentum - the momentum to use when computing the EMA of the model's weights: new_average = - momentum * old_average + (1 - momentum) * current_variable_value - frequency - the number of iterations, to overwrite the model variable by its moving average. - """ - self._kwargs["use_ema"] = enable - if not enable: - logger.debug("ema disabled.") - return - - logger.info("Enabling EMA") - self._kwargs["ema_momentum"] = momentum - self._kwargs["ema_overwrite_frequency"] = frequency - logger.debug("ema enabled (momentum: %s, frequency: %s)", momentum, frequency) - - def _configure_kwargs(self, weight_decay: float, gradient_accumulation_steps: int) -> None: - """Configure the remaining global optimizer kwargs - - Parameters - ---------- - weight_decay - The amount of weight decay to apply - gradient_accumulation_steps - The number of steps to accumulate gradients for before applying the average - """ - if weight_decay > 0.0: - logger.info("Enabling Weight Decay: %s", weight_decay) - self._kwargs["weight_decay"] = weight_decay - else: - logger.debug("weight decay disabled") - - if gradient_accumulation_steps > 1: - logger.info("Enabling Gradient Accumulation: %s", gradient_accumulation_steps) - self._kwargs["gradient_accumulation_steps"] = gradient_accumulation_steps - else: - logger.debug("gradient accumulation disabled") - - def _configure_specific(self) -> None: - """Configure keyword optimizer specific keyword arguments based on user settings.""" - opts = self._valid[cfg_opt.optimizer()][1] - if not opts: - logger.debug("No additional kwargs to set for '%s'", cfg_opt.optimizer()) - return - - for key, val in opts.items(): - opt_val = getattr(cfg_opt, key)() - logger.debug("Setting kwarg '%s' from '%s' to: %s", val, key, opt_val) - self._kwargs[val] = opt_val - - def _configure(self) -> None: - """Process the user configuration options into Keras Optimizer kwargs.""" - self._configure_clipping(T.cast(T.Literal["autoclip", "norm", "value", "none"], - cfg_opt.gradient_clipping()), - cfg_opt.clipping_value(), - cfg_opt.autoclip_history()) - - self._configure_ema(cfg_opt.use_ema(), - cfg_opt.ema_momentum(), - cfg_opt.ema_frequency()) - - self._configure_kwargs(cfg_opt.weight_decay(), - cfg_opt.gradient_accumulation()) - - self._configure_specific() - - logger.debug("Configured '%s' optimizer. kwargs: %s", cfg_opt.optimizer(), self._kwargs) - - class Settings(): """Core training settings. diff --git a/plugins/train/train_config.py b/plugins/train/train_config.py index d4835946f8..5ef196712a 100644 --- a/plugins/train/train_config.py +++ b/plugins/train/train_config.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -""" Default configurations for models """ +"""Default configurations for models""" import gettext import logging @@ -23,10 +23,10 @@ class _Config(FaceswapConfig): - """ Config File for Models """ + """Config File for Models""" # pylint:disable=too-many-statements def set_defaults(self, helptext="") -> None: - """ Set the default values for config """ + """Set the default values for config""" super().set_defaults(helptext=_("Options that apply to all models") + _ADDITIONAL_INFO) self._defaults_from_plugin(os.path.dirname(__file__)) for section, opts in trainer_config.get_defaults().items(): @@ -304,7 +304,7 @@ def set_defaults(self, helptext="") -> None: @dataclass class Loss(GlobalSection): - """ global.loss configuration section + """global.loss configuration section Loss Documentation MAE https://heartbeat.fritz.ai/5-regression-loss-functions-all-machine-learners-should-know-4fb140e9d4b0 MSE https://heartbeat.fritz.ai/5-regression-loss-functions-all-machine-learners-should-know-4fb140e9d4b0 @@ -557,7 +557,7 @@ class Loss(GlobalSection): @dataclass class Optimizer(GlobalSection): - """ global.optimizer configuration section """ + """global.optimizer configuration section""" helptext = (_("Optimizer configuration options\n" "The optimizer applies the output of the loss function to the model.\n") + _ADDITIONAL_INFO) @@ -724,37 +724,6 @@ class Optimizer(GlobalSection): min_max=(1, 100), rounding=1, fixed=False) - use_ema = ConfigItem( - datatype=bool, - default=False, - group=_("exponential moving average"), - info=_( - "Enable exponential moving average (EMA). EMA consists of computing an " - "exponential moving average of the weights of the model (as the weight values " - "change after each training batch), and periodically overwriting the weights " - "with their moving average"), - fixed=True) - ema_momentum = ConfigItem( - datatype=float, - default=0.99, - group=_("exponential moving average"), - info=_( - "Only used if use_ema is enabled. This is the momentum to use when computing " - "the EMA of the model's weights: new_average = ema_momentum * old_average + " - "(1 - ema_momentum) * current_variable_value."), - min_max=(0.0, 1.0), - rounding=4, - fixed=True) - ema_frequency = ConfigItem( - datatype=int, - default=100, - group=_("exponential moving average"), - info=_( - "Only used if use_ema is enabled. Set the number of iterations, to overwrite " - "the model variable by its moving average. "), - min_max=(10, 10000), - rounding=10, - fixed=True) ada_beta_1 = ConfigItem( datatype=float, default=0.9, @@ -793,11 +762,11 @@ class Optimizer(GlobalSection): def load_config(config_file: str | None = None) -> None: - """ Load the Train configuration .ini file + """Load the Train configuration .ini file Parameters ---------- - config_file : str | None, optional + config_file Path to a custom .ini configuration file to load. Default: ``None`` (use default configuration file) """ diff --git a/plugins/train/trainer/base.py b/plugins/train/trainer/base.py index 2888f9332b..6eb7861b86 100644 --- a/plugins/train/trainer/base.py +++ b/plugins/train/trainer/base.py @@ -17,6 +17,7 @@ if T.TYPE_CHECKING: from lib.training.data import BatchMeta from lib.training.loss import LossCollator, BatchLoss + from lib.training.optimizer import Optimizer from plugins.train.model._base import ModelBase logger = logging.getLogger(__name__) @@ -119,6 +120,7 @@ def get_sampler(self) -> type[torch.utils.data.RandomSampler | def train_batch(self, inputs: list[torch.Tensor], targets: list[torch.Tensor], + optimizer: Optimizer, meta: BatchMeta) -> list[BatchLoss]: """Override to run a single forward and backwards pass through the model for a single batch @@ -129,6 +131,8 @@ def train_batch(self, targets List of len (num_outputs) of target images in shape (batch_size, num_inputs, height, width, 3) at all model output sizes as float32 0.0 - 1.0 range + optimizer + The configured Optimizer to use meta The meta information for the batch diff --git a/plugins/train/trainer/original.py b/plugins/train/trainer/original.py index d919c1f221..518dbe466b 100644 --- a/plugins/train/trainer/original.py +++ b/plugins/train/trainer/original.py @@ -5,7 +5,6 @@ import logging import typing as T -from keras import ops import torch from lib.utils import get_module_objects @@ -14,6 +13,7 @@ if T.TYPE_CHECKING: from lib.training.data import BatchMeta from lib.training.loss import BatchLoss + from lib.training.optimizer import Optimizer logger = logging.getLogger(__name__) @@ -61,28 +61,24 @@ def _forward(self, logger.trace("Losses: %s", losses) # type:ignore[attr-defined] return losses - def _backwards_and_apply(self, all_loss: torch.Tensor) -> None: + def _backwards_and_apply(self, loss: list[BatchLoss], optimizer: Optimizer) -> None: """Perform the backwards pass on the model Parameters ---------- - all_loss + loss The loss for each output from the model + optimizer + The configured Optimizer to use """ - total_loss = T.cast(torch.Tensor, - self.model.model.optimizer.scale_loss(ops.sum(all_loss))) - total_loss.backward() - - trainable_weights = self.model.model.trainable_weights[:] - gradients = [v.value.grad for v in trainable_weights] - - # Update weights - with torch.no_grad(): - self.model.model.optimizer.apply(gradients, trainable_weights) + total_loss = T.cast(torch.Tensor, sum(x.total for x in loss)) + optimizer.backward(total_loss) + optimizer.step() def train_batch(self, inputs: list[torch.Tensor], targets: list[torch.Tensor], + optimizer: Optimizer, meta: BatchMeta) -> list[BatchLoss]: """Run a single forward and backwards pass through the model for a single batch @@ -93,6 +89,8 @@ def train_batch(self, targets List of len (num_outputs) of target images in shape (batch_size, num_inputs, height, width, 3) at all model output sizes as float32 0.0 - 1.0 range + optimizer + The configured Optimizer to use meta The meta information for the batch @@ -100,10 +98,8 @@ def train_batch(self, ------- The loss for each input to the model in order (A, B, ...) """ - self.model.model.zero_grad() # TODO move this to optimizer loss = self._forward(inputs, targets, meta) - total_loss = T.cast(torch.Tensor, sum(x.total for x in loss)) - self._backwards_and_apply(total_loss) + self._backwards_and_apply(loss, optimizer) return loss diff --git a/scripts/train.py b/scripts/train.py index 2f96d3d7b6..fbac5e3e8c 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -315,6 +315,7 @@ def _load_trainer(self, model: ModelBase) -> Trainer: snapshot_interval=self._args.snapshot_interval) retval = Trainer(PluginLoader.get_trainer(trainer)(model, config), self._args.preview or self._args.write_image or self._args.redirect_gui, + warmup_steps=self._args.warmup, timelapse_folders=[self._args.timelapse_input_a, self._args.timelapse_input_b], timelapse_output=self._args.timelapse_output) diff --git a/tests/lib/model/optimizers_test.py b/tests/lib/model/optimizers_test.py deleted file mode 100644 index 3f986ace50..0000000000 --- a/tests/lib/model/optimizers_test.py +++ /dev/null @@ -1,59 +0,0 @@ -#!/usr/bin/env python3 -""" Tests for Faceswap Initializers. - -Adapted from Keras tests. -""" -import pytest - -import numpy as np - -from keras import device, layers as kl, optimizers as k_optimizers, Sequential - -from lib.model import optimizers -from lib.utils import get_backend - -from tests.utils import generate_test_data, to_categorical - - -def get_test_data(): - """ Obtain randomized test data for training """ - np.random.seed(1337) - (x_train, y_train), _ = generate_test_data(num_train=1000, - num_test=200, - input_shape=(10,), - classification=True, - num_classes=2) - y_train = to_categorical(y_train) - return x_train, y_train - - -def _test_optimizer(optimizer, target=0.75): - x_train, y_train = get_test_data() - - model = Sequential() - model.add(kl.Input((x_train.shape[1], ))) - model.add(kl.Dense(10)) - model.add(kl.Activation("relu")) - model.add(kl.Dense(y_train.shape[1])) - model.add(kl.Activation("softmax")) - model.compile(loss="categorical_crossentropy", - optimizer=optimizer, - metrics=["accuracy"]) - - history = model.fit(x_train, y_train, epochs=2, batch_size=16, verbose=0) # type:ignore - assert history.history["accuracy"][-1] >= target - config = k_optimizers.serialize(optimizer) - optim = k_optimizers.deserialize(config) - new_config = k_optimizers.serialize(optim) - config["class_name"] = config["class_name"].lower() # type:ignore - new_config["class_name"] = new_config["class_name"].lower() # type:ignore - assert config == new_config - - -# TODO remove the next line that supresses a weird pytest bug when it tears down the tempdir -@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") -@pytest.mark.parametrize("dummy", [None], ids=[get_backend().upper()]) -def test_adabelief(dummy): # pylint:disable=unused-argument - """ Test for custom Adam optimizer """ - with device("cpu"): - _test_optimizer(optimizers.AdaBelief(), target=0.20) diff --git a/tests/lib/training/lr_finder_test.py b/tests/lib/training/lr_finder_test.py deleted file mode 100644 index 3fa8641796..0000000000 --- a/tests/lib/training/lr_finder_test.py +++ /dev/null @@ -1,278 +0,0 @@ -#! /usr/env/bin/python3 -""" Unit tests for Learning Rate Finder. """ - -import pytest -import pytest_mock - -import numpy as np -import torch - -from lib.training.lr_finder import LearningRateFinder -from plugins.train import train_config as cfg - -# pylint:disable=unused-import -from tests.lib.config.helpers import patch_config # noqa:[F401] - -# pylint:disable=protected-access,invalid-name,redefined-outer-name - - -class DummyLoss: # pylint:disable=too-few-public-methods - """Dummy loss return value""" - def __init__(self, value): - self.total = torch.Tensor([value]) - - -@pytest.fixture -def _trainer_mock(patch_config, mocker: pytest_mock.MockFixture): # noqa:[F811] - """ Generate a mocked model and feeder object and patch user config items """ - def _apply_patch(iters=1000, mode="default", strength="default"): - patch_config(cfg, {"lr_finder_iterations": iters}) - patch_config(cfg, {"lr_finder_mode": mode}) - patch_config(cfg, {"lr_finder_strength": strength}) - trainer = mocker.MagicMock() - model = mocker.MagicMock() - model.name = "TestModel" - optimizer = mocker.MagicMock() - trainer._plugin.model = model - trainer._plugin.model.model.optimizer = optimizer - return trainer, model, optimizer - return _apply_patch - - -_STRENGTH_LOOKUP = {"default": 10, "aggressive": 5, "extreme": 2.5} - - -_LR_CONF = ((20, "graph_and_set", "default"), - (500, "set", "aggressive"), - (1000, "graph_and_exit", "extreme")) -_LR_CONF_PARAMS = ("iters", "mode", "strength") - -_LR_CMDS = ((4, 0.98), (8, 0.66), (2, 0.33) - ) -_LR_CMDS_PARAMS = ("stop_factor", "beta") -_LR_CMDS_IDS = [f"stop:{x[0]}|beta:{x[1]}" for x in _LR_CMDS] - - -@pytest.mark.parametrize(_LR_CONF_PARAMS, _LR_CONF) -@pytest.mark.parametrize(_LR_CMDS_PARAMS, _LR_CMDS, ids=_LR_CMDS_IDS) -def test_LearningRateFinder_init(iters, mode, strength, stop_factor, beta, _trainer_mock): - """ Test lib.train.LearingRateFinder.__init__ """ - trainer, model, optimizer = _trainer_mock(iters, mode, strength) - lrf = LearningRateFinder(trainer, stop_factor=stop_factor, beta=beta) - assert lrf._trainer is trainer - assert lrf._model is model - assert lrf._optimizer is optimizer - assert lrf._start_lr == 1e-10 - assert lrf._stop_factor == stop_factor - assert lrf._beta == beta - - -_BATCH_END = ((1, 0.01, 1e-5, 0.5), - (27, 0.01, 1e-5, 1e-6), - (42, 0.001, 1e-5, 0.002),) -_BATCH_END_PARAMS = ("iteration", "loss", "learning_rate", "best") -_BATCH_END_IDS = [f"iter:{x[0]}|loss:{x[1]}|lr:{x[2]}" for x in _BATCH_END] - - -@pytest.mark.parametrize(_LR_CMDS_PARAMS, _LR_CMDS, ids=_LR_CMDS_IDS) -@pytest.mark.parametrize(_BATCH_END_PARAMS, _BATCH_END, ids=_BATCH_END_IDS) -def test_LearningRateFinder_on_batch_end(iteration, - loss, - learning_rate, - best, - stop_factor, - beta, - _trainer_mock, - mocker): - """ Test lib.train.LearingRateFinder._on_batch_end """ - trainer, model, optimizer = _trainer_mock() - lrf = LearningRateFinder(trainer, stop_factor=stop_factor, beta=beta) - optimizer.learning_rate.assign = mocker.MagicMock() - optimizer.learning_rate.numpy = mocker.MagicMock(return_value=learning_rate) - - initial_avg = lrf._loss["avg"] - lrf._loss["best"] = best - lrf._on_batch_end(iteration, loss) - - assert lrf._metrics["learning_rates"][-1] == learning_rate - assert lrf._loss["avg"] == (lrf._beta * initial_avg) + ((1 - lrf._beta) * loss) - assert lrf._metrics["losses"][-1] == lrf._loss["avg"] / (1 - (lrf._beta ** iteration)) - - if iteration > 1 and lrf._metrics["losses"][-1] > lrf._stop_factor * lrf._loss["best"]: - assert model.model.stop_training is True - optimizer.learning_rate.assign.assert_not_called() - return - - if iteration == 1: - assert lrf._loss["best"] == lrf._metrics["losses"][-1] - - assert model.model.stop_training is not True - optimizer.learning_rate.assign.assert_called_with( - learning_rate * lrf._lr_multiplier) - - -@pytest.mark.parametrize(_LR_CONF_PARAMS, _LR_CONF) -def test_LearningRateFinder_train(iters, # pylint:disable=too-many-locals - mode, - strength, - _trainer_mock, - mocker): - """ Test lib.train.LearingRateFinder._train """ - trainer, _, _ = _trainer_mock(iters, mode, strength) - - mock_loss_return = [DummyLoss(np.random.random()) for _ in range(2)] - trainer.train_one_batch = mocker.MagicMock(return_value=mock_loss_return) - - lrf = LearningRateFinder(trainer) - - lrf._on_batch_end = mocker.MagicMock() - lrf._update_description = mocker.MagicMock() - - lrf._train() - - trainer.train_one_batch.assert_called() - assert trainer.train_one_batch.call_count == iters - - train_call_args = [mocker.call(x + 1, sum(y.total for y in mock_loss_return)) - for x in range(iters)] - assert lrf._on_batch_end.call_args_list == train_call_args - - lrf._update_description.assert_called() - assert lrf._update_description.call_count == iters - - # NaN break - mock_loss_return = mock_loss_return = [DummyLoss(np.nan) for _ in range(2)] - trainer.train_one_batch = mocker.MagicMock(return_value=mock_loss_return) - - lrf._train() - - assert trainer.train_one_batch.call_count == 1 # Called once - - assert lrf._update_description.call_count == iters # Not called - assert lrf._on_batch_end.call_count == iters # Not called - - -def test_LearningRateFinder_rebuild_optimizer(_trainer_mock): - """ Test lib.train.LearingRateFinder._rebuild_optimizer """ - trainer, _, _ = _trainer_mock() - lrf = LearningRateFinder(trainer) - - class Dummy: - """ Dummy Optimizer""" - name = "test" - - def get_config(self): - """Dummy get_config""" - return {} - - opt = Dummy() - new_opt = lrf._rebuild_optimizer(opt) - assert isinstance(new_opt, Dummy) and opt is not new_opt - - -@pytest.mark.parametrize(_LR_CONF_PARAMS, _LR_CONF) -@pytest.mark.parametrize("new_lr", (1e-4, 3.5e-5, 9.3e-6)) -def test_LearningRateFinder_reset_model(iters, mode, strength, new_lr, _trainer_mock, mocker): - """ Test lib.train.LearingRateFinder._reset_model """ - trainer, model, optimizer = _trainer_mock(iters, mode, strength) - model.state.add_lr_finder = mocker.MagicMock() - model.state.save = mocker.MagicMock() - model.model.load_weights = mocker.MagicMock() - - old_optimizer = optimizer - new_optimizer = mocker.MagicMock() - - def compile_side_effect(*args, **kwargs): # pylint:disable=unused-argument - """ Side effect for model.compile""" - model.model.optimizer = new_optimizer - - model.model.compile.side_effect = compile_side_effect - - lrf = LearningRateFinder(trainer) - lrf._rebuild_optimizer = mocker.MagicMock() - - lrf._reset_model(1e-5, new_lr) - - model.state.add_lr_finder.assert_called_with(new_lr) - model.state.save.assert_called_once() - - if mode == "graph_and_exit": - lrf._rebuild_optimizer.assert_not_called() - model.model.compile.assert_not_called() - model.model.load_weights.assert_not_called() - assert model.model.optimizer is old_optimizer - new_optimizer.learning_rate.assign.assert_not_called() - else: - lrf._rebuild_optimizer.assert_called_once_with(old_optimizer) - model.model.load_weights.assert_called_once() - model.model.compile.assert_called_once() - assert model.model.optimizer is new_optimizer - new_optimizer.learning_rate.assign.assert_called_once_with(new_lr) - - -_LR_FIND = ( - (True, [0.100, 0.050, 0.025], 0.025, [1e-5, 1e-4, 1e-3], "model_exist"), - (False, [0.100, 0.050, 0.025], 0.025, [1e-5, 1e-4, 1e-3], "no_model"), - (True, [0.100, 0.050, 0.025], 0.025, [1e-5, 1e-4, 1e-10], "low_lr"), - ) -_LR_PARAMS_FIND = ("exists", "losses", "best", "learning_rates") - - -@pytest.mark.parametrize(_LR_PARAMS_FIND, - [x[:-1] for x in _LR_FIND], - ids=[x[-1] for x in _LR_FIND]) -@pytest.mark.parametrize(_LR_CONF_PARAMS, _LR_CONF) -@pytest.mark.parametrize(_LR_CMDS_PARAMS, _LR_CMDS[0:1]) -def test_LearningRateFinder_find(iters, # pylint:disable=too-many-arguments,too-many-positional-arguments # noqa[E501] - mode, - strength, - stop_factor, - beta, - exists, - losses, - best, - learning_rates, - _trainer_mock, - mocker): - """ Test lib.train.LearingRateFinder.find """ - # pylint:disable=too-many-locals - trainer, model, optimizer = _trainer_mock(iters, mode, strength) - model.io.model_exists = exists - model.io.save = mocker.MagicMock() - original_lr = float(np.random.rand()) - optimizer.learning_rate.numpy = mocker.MagicMock(return_value=original_lr) - optimizer.learning_rate.assign = mocker.MagicMock() - mocker.patch("shutil.rmtree") - - lrf = LearningRateFinder(trainer, stop_factor=stop_factor, beta=beta) - - train_mock = mocker.MagicMock() - plot_mock = mocker.MagicMock() - reset_mock = mocker.MagicMock() - lrf._train = train_mock - lrf._plot_loss = plot_mock - lrf._reset_model = reset_mock - - lrf._metrics = {"losses": losses, "learning_rates": learning_rates} - lrf._loss = {"best": best} - - result = lrf.find() - - if exists: - model.io.save_assert_not_called() - else: - model.io.save.assert_called_once() - - optimizer.learning_rate.assign.assert_called_with(lrf._start_lr) - train_mock.assert_called_once() - - new_lr = learning_rates[losses.index(best)] / _STRENGTH_LOOKUP[strength] - if new_lr < 1e-9: - plot_mock.assert_not_called() - reset_mock.assert_not_called() - assert not result - return - - plot_mock.assert_called_once() - reset_mock.assert_called_once_with(original_lr, new_lr) - assert result diff --git a/tests/lib/training/lr_warmup_test.py b/tests/lib/training/lr_warmup_test.py deleted file mode 100644 index c1150c8dae..0000000000 --- a/tests/lib/training/lr_warmup_test.py +++ /dev/null @@ -1,181 +0,0 @@ -#!/usr/bin python3 -""" Pytest unit tests for :mod:`lib.training.lr_warmup` """ - -import pytest -import pytest_mock - -from keras.layers import Input, Dense -from keras.models import Model -from keras.optimizers import SGD - -from lib.training import LearningRateWarmup - - -# pylint:disable=protected-access,redefined-outer-name - - -@pytest.fixture -def model_fixture(): - """ Model fixture for testing LR Warmup """ - inp = Input((4, 4, 3)) - var_x = Dense(8)(inp) - model = Model(inputs=inp, outputs=var_x) - model.compile(optimizer=SGD(), loss="mse") - return model - - -_LR_STEPS = [(1e-5, 100), - (3.4e-6, 250), - (9e-4, 599), - (6e-5, 1000)] -_LR_STEPS_IDS = [f"lr:{x[0]}|steps:{x[1]}" for x in _LR_STEPS] - - -@pytest.mark.parametrize(("target_lr", "steps"), _LR_STEPS, ids=_LR_STEPS_IDS) -def test_init(model_fixture: Model, target_lr: float, steps: int) -> None: - """ Test class initializes correctly """ - instance = LearningRateWarmup(model_fixture, target_lr, steps) - - attrs = ["_model", "_target_lr", "_steps", "_current_lr", "_current_step", "_reporting_points"] - assert all(a in instance.__dict__ for a in attrs) - assert all(a in attrs for a in instance.__dict__) - assert instance._current_lr == 0.0 - assert instance._current_step == 0 - - assert isinstance(instance._model, Model) - assert instance._target_lr == target_lr - assert instance._steps == steps - - assert len(instance._reporting_points) == 11 - assert all(isinstance(x, int) for x in instance._reporting_points) - assert instance._reporting_points == [int(steps * i / 10) for i in range(11)] - - -_NOTATION = [(1e-5, "1.0e-05"), - (3.45489e-6, "3.5e-06"), - (0.0004, "4.0e-04"), - (0.1234, "1.2e-01")] - - -@pytest.mark.parametrize(("value", "expected"), _NOTATION, ids=[x[1] for x in _NOTATION]) -def test_format_notation(value: float, expected: str) -> None: - """ Test floats format to string correctly """ - result = LearningRateWarmup._format_notation(value) - assert result == expected - - -_LR_STEPS_CURRENT = [(1e-5, 100, 79), - (3.4e-6, 250, 250), - (9e-4, 599, 0), - (6e-5, 1000, 12)] -_LR_STEPS_CURRENT_IDS = [f"lr:{x[0]}|steps:{x[1]}|current_step:{x[2]}" for x in _LR_STEPS_CURRENT] - - -@pytest.mark.parametrize(("target_lr", "steps", "current_step"), - _LR_STEPS_CURRENT, - ids=_LR_STEPS_CURRENT_IDS) -def test_set_current_learning_rate(model_fixture: Model, - target_lr: float, - steps: int, - current_step: int) -> None: - """ Test that learning rate is set correctly """ - instance = LearningRateWarmup(model_fixture, target_lr, steps) - instance._current_step = current_step - instance._set_learning_rate() - - assert instance._current_lr == instance._current_step / instance._steps * instance._target_lr - assert instance._model.optimizer.learning_rate.value.cpu().numpy() == instance._current_lr - - -_STEPS_CURRENT = [(1000, 1, "start"), - (250, 250, "end"), - (500, 69, "unreported"), - (1000, 200, "reported")] -_STEPS_CURRENT_ID = [f"steps:{x[0]}|current_step:{x[1]}|action:{x[2]}" for x in _STEPS_CURRENT] - - -@pytest.mark.parametrize(("steps", "current_step", "action"), - _STEPS_CURRENT, - ids=_STEPS_CURRENT_ID) -def test_output_status(model_fixture: Model, - steps: int, - current_step: int, - action: str, - mocker: pytest_mock.MockerFixture) -> None: - """ Test that information is output correctly """ - mock_logger = mocker.patch("lib.training.lr_warmup.logger.info") - mock_print = mocker.patch("builtins.print") - instance = LearningRateWarmup(model_fixture, 5e-5, steps) - instance._current_step = current_step - instance._format_notation = mocker.MagicMock() # type:ignore[method-assign] - - instance._output_status() - - if action == "unreported": - assert current_step not in instance._reporting_points - mock_logger.assert_not_called() - instance._format_notation.assert_not_called() # type:ignore[attr-defined] - mock_print.assert_not_called() - return - - mock_logger.assert_called_once() - log_message: str = mock_logger.call_args.args[0] - assert log_message.startswith("[Learning Rate Warmup] ") - - instance._format_notation.assert_called() # type:ignore[attr-defined] - notation_args = [ - x.args for x in instance._format_notation.call_args_list] # type:ignore[attr-defined] - assert all(len(a) == 1 for a in notation_args) - assert all(isinstance(a[0], float) for a in notation_args) - - if action == "start": - mock_print.assert_not_called() - assert all(x in log_message for x in ("Start: ", "Target: ", "Steps: ")) - assert instance._format_notation.call_count == 2 # type:ignore[attr-defined] - return - - if action == "end": - mock_print.assert_called() - assert "Final Learning Rate: " in log_message - instance._format_notation.assert_called_once() # type:ignore[attr-defined] - return - - if action == "reported": - mock_print.assert_called() - assert current_step in instance._reporting_points - assert all(x in log_message for x in ("Step: ", "Current: ", "Target: ")) - assert instance._format_notation.call_count == 2 # type:ignore[attr-defined] - - -_STEPS_CURRENT_CALL = [(0, 500, "disabled"), - (1000, 500, "progress"), - (1000, 1000, "completed"), - (1000, 1111, "completed2")] -_STEPS_CURRENT_CALL_ID = [f"steps:{x[0]}|current_step:{x[1]}|action:{x[2]}" - for x in _STEPS_CURRENT_CALL] - - -@pytest.mark.parametrize(("steps", "current_step", "action"), - _STEPS_CURRENT_CALL, - ids=_STEPS_CURRENT_CALL_ID) -def test__call__(model_fixture: Model, - steps: int, - current_step: int, - action: str, - mocker: pytest_mock.MockerFixture) -> None: - """ Test calling the instance works correctly """ - instance = LearningRateWarmup(model_fixture, 5e-5, steps) - instance._current_step = current_step - instance._set_learning_rate = mocker.MagicMock() # type:ignore[method-assign] - instance._output_status = mocker.MagicMock() # type:ignore[method-assign] - - instance() - - if action in ("disabled", "completed", "completed2"): - assert instance._current_step == current_step - instance._set_learning_rate.assert_not_called() # type:ignore[attr-defined] - instance._output_status.assert_not_called() # type:ignore[attr-defined] - else: - assert instance._current_step == current_step + 1 - instance._set_learning_rate.assert_called_once() # type:ignore[attr-defined] - instance._output_status.assert_called_once() # type:ignore[attr-defined] diff --git a/tests/plugins/train/trainer/test_original.py b/tests/plugins/train/trainer/test_original.py index a386f6fb1e..41a30a1ec6 100644 --- a/tests/plugins/train/trainer/test_original.py +++ b/tests/plugins/train/trainer/test_original.py @@ -14,7 +14,7 @@ class DummyLoss: # pylint:disable=too-few-public-methods """Dummy loss return""" - total = 1.0 + total = np.random.rand() @pytest.fixture @@ -52,12 +52,11 @@ def test_Trainer_train_batch(_trainer_mocked, mocker): instance._backwards_and_apply = mocker.MagicMock() instance.model.model.zero_grad = mocker.MagicMock() - ret_val = instance.train_batch("TEST_INPUT", "TEST_TARGET", "TEST_META") + ret_val = instance.train_batch("TEST_INPUT", "TEST_TARGET", "TEST_OPTIMIZER", "TEST_META") assert ret_val == loss_return instance._forward.assert_called_once_with("TEST_INPUT", "TEST_TARGET", "TEST_META") - instance._backwards_and_apply.assert_called_once_with(1.0) - instance.model.model.zero_grad.assert_called_once() + instance._backwards_and_apply.assert_called_once_with(loss_return, "TEST_OPTIMIZER") @pytest.mark.parametrize("outputs", (1, 2, 4)) @@ -114,19 +113,9 @@ def test_Trainer_backwards_and_apply(_trainer_mocked, mocker): """ Test that original trainer _backwards_and_apply calls the correct model methods """ instance = _trainer_mocked() - mock_loss = mocker.MagicMock() - instance.model.model.optimizer.scale_loss = mocker.MagicMock(return_value=mock_loss) - instance.model.model.optimizer.app = mocker.MagicMock(return_value=mock_loss) + mock_optimizer = mocker.MagicMock() + all_loss = [DummyLoss] + instance._backwards_and_apply(all_loss, mock_optimizer) - all_loss = np.random.rand() - instance._backwards_and_apply(all_loss) - - scale_mock = instance.model.model.optimizer.scale_loss - scale_mock.assert_called_once() - assert not scale_mock.call_args[1] - assert len(scale_mock.call_args[0]) == 1 - assert np.isclose(all_loss, scale_mock.call_args[0][0].cpu().numpy()) - - mock_loss.backward.assert_called_once() - - instance.model.model.optimizer.apply.assert_called_once() + mock_optimizer.backward.assert_called_once_with(all_loss[0].total) + mock_optimizer.step.assert_called_once()