130 lines
4.0 KiB
Python
130 lines
4.0 KiB
Python
import logging
|
|
from itertools import islice
|
|
from typing import Collection, Optional, Tuple
|
|
|
|
from torch import nn as nn
|
|
|
|
from timm.models import group_parameters
|
|
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
def param_groups_weight_decay(
|
|
model: nn.Module,
|
|
weight_decay: float = 1e-5,
|
|
no_weight_decay_list: Collection[str] = (),
|
|
):
|
|
no_weight_decay_list = set(no_weight_decay_list)
|
|
decay = []
|
|
no_decay = []
|
|
for name, param in model.named_parameters():
|
|
if not param.requires_grad:
|
|
continue
|
|
|
|
if param.ndim <= 1 or name.endswith(".bias") or name in no_weight_decay_list:
|
|
no_decay.append(param)
|
|
else:
|
|
decay.append(param)
|
|
|
|
return [
|
|
{'params': no_decay, 'weight_decay': 0.},
|
|
{'params': decay, 'weight_decay': weight_decay}]
|
|
|
|
|
|
def _group(it, size):
|
|
it = iter(it)
|
|
return iter(lambda: tuple(islice(it, size)), ())
|
|
|
|
|
|
def _layer_map(model, layers_per_group=12, num_groups=None):
|
|
def _in_head(n, hp):
|
|
if not hp:
|
|
return True
|
|
elif isinstance(hp, (tuple, list)):
|
|
return any([n.startswith(hpi) for hpi in hp])
|
|
else:
|
|
return n.startswith(hp)
|
|
|
|
head_prefix = getattr(model, 'pretrained_cfg', {}).get('classifier', None)
|
|
names_trunk = []
|
|
names_head = []
|
|
for n, _ in model.named_parameters():
|
|
names_head.append(n) if _in_head(n, head_prefix) else names_trunk.append(n)
|
|
|
|
# group non-head layers
|
|
num_trunk_layers = len(names_trunk)
|
|
if num_groups is not None:
|
|
layers_per_group = -(num_trunk_layers // -num_groups)
|
|
names_trunk = list(_group(names_trunk, layers_per_group))
|
|
|
|
num_trunk_groups = len(names_trunk)
|
|
layer_map = {n: i for i, l in enumerate(names_trunk) for n in l}
|
|
layer_map.update({n: num_trunk_groups for n in names_head})
|
|
return layer_map
|
|
|
|
|
|
def param_groups_layer_decay(
|
|
model: nn.Module,
|
|
weight_decay: float = 0.05,
|
|
no_weight_decay_list: Collection[str] = (),
|
|
weight_decay_exclude_1d: bool = True,
|
|
layer_decay: float = .75,
|
|
end_layer_decay: Optional[float] = None,
|
|
verbose: bool = False,
|
|
):
|
|
"""
|
|
Parameter groups for layer-wise lr decay & weight decay
|
|
Based on BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
|
|
"""
|
|
no_weight_decay_list = set(no_weight_decay_list)
|
|
param_group_names = {} # NOTE for debugging
|
|
param_groups = {}
|
|
|
|
if hasattr(model, 'group_matcher'):
|
|
# FIXME interface needs more work
|
|
layer_map = group_parameters(model, model.group_matcher(coarse=False), reverse=True)
|
|
else:
|
|
# fallback
|
|
layer_map = _layer_map(model)
|
|
num_layers = max(layer_map.values()) + 1
|
|
layer_max = num_layers - 1
|
|
layer_scales = list(layer_decay ** (layer_max - i) for i in range(num_layers))
|
|
|
|
for name, param in model.named_parameters():
|
|
if not param.requires_grad:
|
|
continue
|
|
|
|
# no decay: all 1D parameters and model specific ones
|
|
if (weight_decay_exclude_1d and param.ndim <= 1) or name in no_weight_decay_list:
|
|
g_decay = "no_decay"
|
|
this_decay = 0.
|
|
else:
|
|
g_decay = "decay"
|
|
this_decay = weight_decay
|
|
|
|
layer_id = layer_map.get(name, layer_max)
|
|
group_name = "layer_%d_%s" % (layer_id, g_decay)
|
|
|
|
if group_name not in param_groups:
|
|
this_scale = layer_scales[layer_id]
|
|
param_group_names[group_name] = {
|
|
"lr_scale": this_scale,
|
|
"weight_decay": this_decay,
|
|
"param_names": [],
|
|
}
|
|
param_groups[group_name] = {
|
|
"lr_scale": this_scale,
|
|
"weight_decay": this_decay,
|
|
"params": [],
|
|
}
|
|
|
|
param_group_names[group_name]["param_names"].append(name)
|
|
param_groups[group_name]["params"].append(param)
|
|
|
|
if verbose:
|
|
import json
|
|
_logger.info("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))
|
|
|
|
return list(param_groups.values())
|