diff --git a/timm/optim/optim_factory.py b/timm/optim/optim_factory.py index 06788d2e..c82fd3d2 100644 --- a/timm/optim/optim_factory.py +++ b/timm/optim/optim_factory.py @@ -1,7 +1,7 @@ """ Optimizer Factory w/ Custom Weight Decay Hacked together by / Copyright 2021 Ross Wightman """ -import json +import logging from itertools import islice from typing import Optional, Callable, Tuple @@ -31,6 +31,8 @@ try: except ImportError: has_apex = False +_logger = logging.getLogger(__name__) + def param_groups_weight_decay( model: nn.Module, @@ -92,6 +94,7 @@ def param_groups_layer_decay( no_weight_decay_list: Tuple[str] = (), layer_decay: float = .75, end_layer_decay: Optional[float] = None, + verbose: bool = False, ): """ Parameter groups for layer-wise lr decay & weight decay @@ -142,8 +145,9 @@ def param_groups_layer_decay( param_group_names[group_name]["param_names"].append(name) param_groups[group_name]["params"].append(param) - # FIXME temporary output to debug new feature - print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2)) + if verbose: + import json + _logger.info("parameter groups: \n%s" % json.dumps(param_group_names, indent=2)) return list(param_groups.values())