|
|
|
|
|
|
|
|
|
|
|
|
| import fnmatch
|
| import inspect
|
| import itertools
|
| import logging
|
| import types
|
| from typing import (
|
| Any,
|
| Callable,
|
| Dict,
|
| Iterable,
|
| List,
|
| Mapping,
|
| Optional,
|
| Set,
|
| Tuple,
|
| Type,
|
| Union,
|
| )
|
|
|
| import hydra
|
|
|
| import torch
|
| import torch.nn as nn
|
| from omegaconf import DictConfig
|
| from torch import Tensor
|
|
|
|
|
| class Optimizer:
|
| def __init__(self, optimizer, schedulers=None) -> None:
|
| self.optimizer = optimizer
|
| self.schedulers = schedulers
|
| self._validate_optimizer_schedulers()
|
| self.step_schedulers(0.0, 0)
|
|
|
| def _validate_optimizer_schedulers(self):
|
| if self.schedulers is None:
|
| return
|
| for _, set_of_schedulers in enumerate(self.schedulers):
|
| for option, _ in set_of_schedulers.items():
|
| assert option in self.optimizer.defaults, (
|
| "Optimizer option "
|
| f"{option} not found in {self.optimizer}. Valid options are "
|
| f"{self.optimizer.defaults.keys()}"
|
| )
|
|
|
| def step_schedulers(self, where: float, step: int) -> None:
|
| if self.schedulers is None:
|
| return
|
| for i, param_group in enumerate(self.optimizer.param_groups):
|
| for option, scheduler in self.schedulers[i].items():
|
| if "step" in inspect.signature(scheduler.__call__).parameters:
|
| new_value = scheduler(step=step, where=where)
|
| elif (
|
| hasattr(scheduler, "scheduler")
|
| and "step"
|
| in inspect.signature(scheduler.scheduler.__call__).parameters
|
| ):
|
|
|
| new_value = scheduler(step=step, where=where)
|
| else:
|
| new_value = scheduler(where)
|
| param_group[option] = new_value
|
|
|
| def step(self, where, step, closure=None):
|
| self.step_schedulers(where, step)
|
| return self.optimizer.step(closure)
|
|
|
| def zero_grad(self, *args, **kwargs):
|
| return self.optimizer.zero_grad(*args, **kwargs)
|
|
|
|
|
| def set_default_parameters(
|
| scheduler_cfgs: List[DictConfig], all_parameter_names: Set[str]
|
| ) -> None:
|
| """Set up the "default" scheduler with the right parameters.
|
|
|
| Args:
|
| scheduler_cgfs: A list of scheduler configs, where each scheduler also
|
| specifies which parameters it applies to, based on the names of parameters
|
| or the class of the modules. At most one scheduler is allowed to skip this
|
| specification, which is used as a "default" specification for any remaining
|
| parameters.
|
| all_parameter_names: Names of all the parameters to consider.
|
| """
|
| constraints = [
|
| scheduler_cfg.parameter_names
|
| for scheduler_cfg in scheduler_cfgs
|
| if scheduler_cfg.parameter_names is not None
|
| ]
|
| if len(constraints) == 0:
|
| default_params = set(all_parameter_names)
|
| else:
|
| default_params = all_parameter_names - set.union(*constraints)
|
| default_count = 0
|
| for scheduler_cfg in scheduler_cfgs:
|
| if scheduler_cfg.parameter_names is None:
|
| scheduler_cfg.parameter_names = default_params
|
| default_count += 1
|
| assert default_count <= 1, "Only one scheduler per option can be default"
|
| if default_count == 0:
|
|
|
|
|
| scheduler_cfgs.append({"parameter_names": default_params})
|
|
|
|
|
| def name_constraints_to_parameters(
|
| param_constraints: List[Set[str]], named_parameters: Dict[str, Tensor]
|
| ) -> List[torch.nn.Parameter]:
|
| """Return parameters which match the intersection of parameter constraints.
|
|
|
| Note that this returns the parameters themselves, not their names.
|
|
|
| Args:
|
| param_constraints: A list, with each element being a set of allowed parameters.
|
| named_parameters: Mapping from a parameter name to the parameter itself.
|
|
|
| Returns:
|
| A list containing the parameters which overlap with _each_ constraint set from
|
| param_constraints.
|
| """
|
| matching_names = set.intersection(*param_constraints)
|
| return [value for name, value in named_parameters.items() if name in matching_names]
|
|
|
|
|
| def map_scheduler_cfgs_to_param_groups(
|
| all_scheduler_cfgs: Iterable[List[Dict]],
|
| named_parameters: Dict[str, Tensor],
|
| ) -> Tuple[List[Dict[Any, Any]], List[Dict[str, List[torch.nn.Parameter]]]]:
|
| """Produce parameter groups corresponding to all the scheduler configs.
|
|
|
| Takes all the scheduler configs, each of which applies to a specific optimizer
|
| option (like "lr" or "weight_decay") and has a set of parameter names which it
|
| applies to, and produces a final set of param groups where each param group
|
| covers all the options which apply to a particular set of parameters.
|
|
|
| Args:
|
| all_scheduler_cfgs: All the scheduler configs covering every option.
|
| named_parameters: Mapping from a parameter name to the parameter itself.
|
| Returns:
|
| Tuple of lists of schedulers and param_groups, where schedulers[i]
|
| applies to param_groups[i].
|
| """
|
|
|
| scheduler_cfgs_per_param_group = itertools.product(*all_scheduler_cfgs)
|
| schedulers = []
|
| param_groups = []
|
| for scheduler_cfgs in scheduler_cfgs_per_param_group:
|
| param_constraints = [
|
| scheduler_cfg["parameter_names"] for scheduler_cfg in scheduler_cfgs
|
| ]
|
| matching_parameters = name_constraints_to_parameters(
|
| param_constraints, named_parameters
|
| )
|
| if len(matching_parameters) == 0:
|
| continue
|
| schedulers_for_group = {
|
| scheduler_cfg["option"]: scheduler_cfg["scheduler"]
|
| for scheduler_cfg in scheduler_cfgs
|
| if "option" in scheduler_cfg
|
| }
|
| schedulers.append(schedulers_for_group)
|
| param_groups.append({"params": matching_parameters})
|
| return schedulers, param_groups
|
|
|
|
|
| def validate_param_group_params(param_groups: List[Dict], model: nn.Module):
|
| """Check that the param groups are non-overlapping and cover all the parameters.
|
|
|
| Args:
|
| param_groups: List of all param groups
|
| model: Model to validate against. The check ensures that all the model
|
| parameters are part of param_groups
|
| """
|
| for pg in param_groups:
|
|
|
| assert len(pg["params"]) == len(set(pg["params"]))
|
| parameters = [set(param_group["params"]) for param_group in param_groups]
|
| model_parameters = {parameter for _, parameter in model.named_parameters()}
|
| for p1, p2 in itertools.permutations(parameters, 2):
|
| assert p1.isdisjoint(p2), "Scheduler generated param_groups should be disjoint"
|
| assert set.union(*parameters) == model_parameters, (
|
| "Scheduler generated param_groups must include all parameters of the model."
|
| f" Found {len(set.union(*parameters))} params whereas model has"
|
| f" {len(model_parameters)} params"
|
| )
|
|
|
|
|
| def unix_module_cls_pattern_to_parameter_names(
|
| filter_module_cls_names: List[str],
|
| module_cls_to_param_names: Dict[Type, str],
|
| ) -> Union[None, Set[str]]:
|
| """Returns param names which pass the filters specified in filter_module_cls_names.
|
|
|
| Args:
|
| filter_module_cls_names: A list of filter strings containing class names, like
|
| ["torch.nn.LayerNorm", "torch.nn.BatchNorm2d"]
|
| module_cls_to_param_names: Mapping from module classes to the parameter names
|
| they contain. See `get_module_cls_to_param_names`.
|
| """
|
| if filter_module_cls_names is None:
|
| return set()
|
| allowed_parameter_names = []
|
| for module_cls_name in filter_module_cls_names:
|
| module_cls = hydra.utils.get_class(module_cls_name)
|
| if module_cls not in module_cls_to_param_names:
|
| raise AssertionError(
|
| f"module_cls_name {module_cls_name} does not "
|
| "match any classes in the model"
|
| )
|
| matching_parameters = module_cls_to_param_names[module_cls]
|
| assert (
|
| len(matching_parameters) > 0
|
| ), f"module_cls_name {module_cls_name} does not contain any parameters in the model"
|
| logging.info(
|
| f"Matches for module_cls_name [{module_cls_name}]: {matching_parameters} "
|
| )
|
| allowed_parameter_names.append(matching_parameters)
|
| return set.union(*allowed_parameter_names)
|
|
|
|
|
| def unix_param_pattern_to_parameter_names(
|
| filter_param_names: Optional[List[str]],
|
| parameter_names: Dict[str, torch.Tensor],
|
| ) -> Union[None, Set[str]]:
|
| """Returns param names which pass the filters specified in filter_param_names.
|
|
|
| Args:
|
| filter_param_names: A list of unix-style filter strings with optional
|
| wildcards, like ["block.2.*", "block.2.linear.weight"]
|
| module_cls_to_param_names: Mapping from module classes to the parameter names
|
| they contain. See `get_module_cls_to_param_names`.
|
| """
|
|
|
| if filter_param_names is None:
|
| return set()
|
| allowed_parameter_names = []
|
| for param_name in filter_param_names:
|
| matching_parameters = set(fnmatch.filter(parameter_names, param_name))
|
| assert (
|
| len(matching_parameters) >= 1
|
| ), f"param_name {param_name} does not match any parameters in the model"
|
| logging.info(f"Matches for param_name [{param_name}]: {matching_parameters}")
|
| allowed_parameter_names.append(matching_parameters)
|
| return set.union(*allowed_parameter_names)
|
|
|
|
|
| def _unix_pattern_to_parameter_names(
|
| scheduler_cfg: DictConfig,
|
| parameter_names: Set[str],
|
| module_cls_to_param_names: Dict[Type, str],
|
| ) -> Union[None, Set[str]]:
|
| """Returns param names which pass the filters specified in scheduler_cfg.
|
|
|
| Args:
|
| scheduler_cfg: The config for the scheduler
|
| parameter_names: The set of all parameter names which will be filtered
|
| """
|
| if "param_names" not in scheduler_cfg and "module_cls_names" not in scheduler_cfg:
|
| return None
|
| return unix_param_pattern_to_parameter_names(
|
| scheduler_cfg.get("param_names"), parameter_names
|
| ).union(
|
| unix_module_cls_pattern_to_parameter_names(
|
| scheduler_cfg.get("module_cls_names"), module_cls_to_param_names
|
| )
|
| )
|
|
|
|
|
| def get_module_cls_to_param_names(
|
| model: nn.Module, param_allowlist: Set[str] = None
|
| ) -> Dict[Type, str]:
|
| """Produce a mapping from all the modules classes to the names of parames they own.
|
|
|
| Only counts a parameter as part of the immediate parent module, i.e. recursive
|
| parents do not count.
|
|
|
| Args:
|
| model: Model to iterate over
|
| param_allowlist: If specified, only these param names will be processed
|
| """
|
|
|
| module_cls_to_params = {}
|
| for module_name, module in model.named_modules():
|
| module_cls = type(module)
|
| module_cls_to_params.setdefault(module_cls, set())
|
| for param_name, _ in module.named_parameters(recurse=False):
|
| full_param_name = get_full_parameter_name(module_name, param_name)
|
| if param_allowlist is None or full_param_name in param_allowlist:
|
| module_cls_to_params[module_cls].add(full_param_name)
|
| return module_cls_to_params
|
|
|
|
|
| def construct_optimizer(
|
| model: torch.nn.Module,
|
| optimizer_conf: Any,
|
| options_conf: Mapping[str, List] = None,
|
| param_group_modifiers_conf: List[Callable] = None,
|
| param_allowlist: Optional[Set[str]] = None,
|
| validate_param_groups=True,
|
| ) -> Optimizer:
|
| """
|
| Constructs a stochastic gradient descent or ADAM (or ADAMw) optimizer
|
| with momentum. i.e, constructs a torch.optim.Optimizer with zero-weight decay
|
| Batchnorm and/or no-update 1-D parameters support, based on the config.
|
|
|
| Supports wrapping the optimizer with Layer-wise Adaptive Rate Scaling
|
| (LARS): https://arxiv.org/abs/1708.03888
|
|
|
| Args:
|
| model: model to perform stochastic gradient descent
|
| optimization or ADAM optimization.
|
| optimizer_conf: Hydra config consisting a partial torch optimizer like SGD or
|
| ADAM, still missing the params argument which this function provides to
|
| produce the final optimizer
|
| param_group_modifiers_conf: Optional user specified functions which can modify
|
| the final scheduler configs before the optimizer's param groups are built
|
| param_allowlist: The parameters to optimize. Parameters which are not part of
|
| this allowlist will be skipped.
|
| validate_param_groups: If enabled, valides that the produced param_groups don't
|
| overlap and cover all the model parameters.
|
| """
|
| if param_allowlist is None:
|
| param_allowlist = {name for name, _ in model.named_parameters()}
|
|
|
| named_parameters = {
|
| name: param
|
| for name, param in model.named_parameters()
|
| if name in param_allowlist
|
| }
|
|
|
| if not options_conf:
|
| optimizer = hydra.utils.instantiate(optimizer_conf, named_parameters.values())
|
| return Optimizer(optimizer)
|
|
|
| all_parameter_names = {
|
| name for name, _ in model.named_parameters() if name in param_allowlist
|
| }
|
| module_cls_to_all_param_names = get_module_cls_to_param_names(
|
| model, param_allowlist
|
| )
|
|
|
| scheduler_cfgs_per_option = hydra.utils.instantiate(options_conf)
|
| all_scheduler_cfgs = []
|
| for option, scheduler_cfgs in scheduler_cfgs_per_option.items():
|
| for config in scheduler_cfgs:
|
| config.option = option
|
| config.parameter_names = _unix_pattern_to_parameter_names(
|
| config, all_parameter_names, module_cls_to_all_param_names
|
| )
|
| set_default_parameters(scheduler_cfgs, all_parameter_names)
|
| all_scheduler_cfgs.append(scheduler_cfgs)
|
|
|
| if param_group_modifiers_conf:
|
| for custom_param_modifier in param_group_modifiers_conf:
|
| custom_param_modifier = hydra.utils.instantiate(custom_param_modifier)
|
| all_scheduler_cfgs = custom_param_modifier(
|
| scheduler_cfgs=all_scheduler_cfgs, model=model
|
| )
|
| schedulers, param_groups = map_scheduler_cfgs_to_param_groups(
|
| all_scheduler_cfgs, named_parameters
|
| )
|
| if validate_param_groups:
|
| validate_param_group_params(param_groups, model)
|
| optimizer = hydra.utils.instantiate(optimizer_conf, param_groups)
|
| return Optimizer(optimizer, schedulers)
|
|
|
|
|
| def get_full_parameter_name(module_name, param_name):
|
| if module_name == "":
|
| return param_name
|
| return f"{module_name}.{param_name}"
|
|
|
|
|
| class GradientClipper:
|
| """
|
| Gradient clipping utils that works for DDP
|
| """
|
|
|
| def __init__(self, max_norm: float = 1.0, norm_type: int = 2):
|
| assert isinstance(max_norm, (int, float)) or max_norm is None
|
| self.max_norm = max_norm if max_norm is None else float(max_norm)
|
| self.norm_type = norm_type
|
|
|
| def __call__(self, model: nn.Module):
|
| if self.max_norm is None:
|
| return
|
|
|
| nn.utils.clip_grad_norm_(
|
| model.parameters(), max_norm=self.max_norm, norm_type=self.norm_type
|
| )
|
|
|
|
|
| class ValueScaler:
|
| def __init__(self, scheduler, mult_val: float):
|
| self.scheduler = scheduler
|
| self.mult_val = mult_val
|
|
|
| def __call__(self, *args, **kwargs):
|
| val = self.scheduler(*args, **kwargs)
|
| return val * self.mult_val
|
|
|
|
|
| def rgetattr(obj, rattrs: str = None):
|
| """
|
| Like getattr(), but supports dotted notation for nested objects.
|
| rattrs is a str of form 'attr1.attr2', returns obj.attr1.attr2
|
| """
|
| if rattrs is None:
|
| return obj
|
| attrs = rattrs.split(".")
|
| for attr in attrs:
|
| obj = getattr(obj, attr)
|
| return obj
|
|
|
|
|
| def layer_decay_param_modifier(
|
| scheduler_cfgs: List[List[Dict]],
|
| model,
|
| layer_decay_value: float,
|
| layer_decay_min: Optional[float] = None,
|
| apply_to: Optional[str] = None,
|
| overrides: List[Dict] = (),
|
| ) -> List[List[Dict]]:
|
| """
|
| Args
|
| - scheduler_cfgs: a list of omegaconf.ListConfigs.
|
| Each element in the list is a omegaconfg.DictConfig with the following structure
|
| {
|
| "scheduler": <some fvcore scheduler>
|
| "option": <value> possible options are "lr", "weight_decay" etc.
|
| "parameter_names": Set of str indicating param names that this scheduler applies to
|
| }
|
| - model: a model that implements a method `get_layer_id` that maps layer_name to an integer and
|
| and a method get_num_layers.
|
| Alternatively, use apply_to argument to select a specific component of the model.
|
| - layer_decay_value: float
|
| - layer_decay_min: min val for layer decay
|
| - apply_to: optional arg to select which component of the model to apply the the layer decay modifier to
|
| - overrides: to manually override lr for specific patterns. Is a list of dicts. Each dict, has keys "pattern", "value".
|
| Returns
|
| - scheduler_configs: same structure as the input, elements can be modified
|
| """
|
| model = rgetattr(model, apply_to)
|
| num_layers = model.get_num_layers() + 1
|
| layer_decays = [
|
| layer_decay_value ** (num_layers - i) for i in range(num_layers + 1)
|
| ]
|
| if layer_decay_min is not None:
|
| layer_decays = [max(val, layer_decay_min) for val in layer_decays]
|
| final_scheduler_cfgs = []
|
|
|
| for scheduler_cfg_group in scheduler_cfgs:
|
| curr_cfg_group = []
|
|
|
| for scheduler_cfg in scheduler_cfg_group:
|
| if scheduler_cfg["option"] != "lr":
|
| curr_cfg_group.append(scheduler_cfg)
|
| continue
|
|
|
|
|
|
|
| parameter_names = sorted(scheduler_cfg["parameter_names"])
|
|
|
|
|
| layer_cfg_groups = {}
|
| for param_name in parameter_names:
|
| layer_id = num_layers
|
| this_scale = layer_decays[layer_id]
|
| if param_name.startswith(apply_to):
|
| layer_id = model.get_layer_id(param_name)
|
| this_scale = layer_decays[layer_id]
|
|
|
| for override in overrides:
|
| if fnmatch.fnmatchcase(param_name, override["pattern"]):
|
| this_scale = float(override["value"])
|
| layer_id = override["pattern"]
|
| break
|
|
|
| if layer_id not in layer_cfg_groups:
|
| curr_param = {
|
| "option": scheduler_cfg["option"],
|
| "scheduler": ValueScaler(
|
| scheduler_cfg["scheduler"], this_scale
|
| ),
|
| "parameter_names": {param_name},
|
| }
|
| else:
|
| curr_param = layer_cfg_groups[layer_id]
|
| curr_param["parameter_names"].add(param_name)
|
| layer_cfg_groups[layer_id] = curr_param
|
|
|
| for layer_cfg in layer_cfg_groups.values():
|
| curr_cfg_group.append(layer_cfg)
|
|
|
| final_scheduler_cfgs.append(curr_cfg_group)
|
| return final_scheduler_cfgs
|
|
|
