503 lines
19 KiB
Python
503 lines
19 KiB
Python
|
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||
|
|
# All rights reserved.
|
||
|
|
|
||
|
|
# This source code is licensed under the license found in the
|
||
|
|
# LICENSE file in the root directory of this source tree.
|
||
|
|
|
||
|
|
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
|
||
|
|
):
|
||
|
|
# To handle ValueScaler wrappers
|
||
|
|
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:
|
||
|
|
# No default scheduler specified, add a default, but without any scheduler
|
||
|
|
# for that option
|
||
|
|
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: # If no overlap of parameters, skip
|
||
|
|
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:
|
||
|
|
# no param should be repeated within a group
|
||
|
|
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 # no-op
|
||
|
|
|
||
|
|
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 = []
|
||
|
|
# scheduler_cfgs is a list of lists
|
||
|
|
for scheduler_cfg_group in scheduler_cfgs:
|
||
|
|
curr_cfg_group = []
|
||
|
|
# scheduler_cfg_group is a list of dictionaries
|
||
|
|
for scheduler_cfg in scheduler_cfg_group:
|
||
|
|
if scheduler_cfg["option"] != "lr":
|
||
|
|
curr_cfg_group.append(scheduler_cfg)
|
||
|
|
continue
|
||
|
|
# Need sorted so that the list of parameter names is deterministic and consistent
|
||
|
|
# across re-runs of this job. Else it was causing issues with loading the optimizer
|
||
|
|
# state during a job restart (D38591759)
|
||
|
|
parameter_names = sorted(scheduler_cfg["parameter_names"])
|
||
|
|
|
||
|
|
# Only want one cfg group per layer
|
||
|
|
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]
|
||
|
|
# Overrides
|
||
|
|
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
|