From 690f31d02d3732767c5a68194e5fdaa52c790fad Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Sat, 23 Oct 2021 20:30:53 -0700 Subject: [PATCH] Post merge cleanup, restore previous unwrap fn --- timm/utils/model.py | 56 ++++++++++++++++++++++++++------------------- 1 file changed, 32 insertions(+), 24 deletions(-) diff --git a/timm/utils/model.py b/timm/utils/model.py index b9f3e9d3..3fbc9f25 100644 --- a/timm/utils/model.py +++ b/timm/utils/model.py @@ -7,33 +7,38 @@ import fnmatch import torch from torchvision.ops.misc import FrozenBatchNorm2d -from .model_ema import ModelEma +_SUB_MODULE_ATTR = ('module', 'model') -def unwrap_model(model): - if isinstance(model, ModelEma): - return unwrap_model(model.ema) - else: - return model.module if hasattr(model, 'module') else model + +def unwrap_model(model, recursive=True): + for attr in _SUB_MODULE_ATTR: + sub_module = getattr(model, attr, None) + if sub_module is not None: + return unwrap_model(sub_module) if recursive else sub_module + return model def get_state_dict(model, unwrap_fn=unwrap_model): return unwrap_fn(model).state_dict() -def avg_sq_ch_mean(model, input, output): - "calculate average channel square mean of output activations" - return torch.mean(output.mean(axis=[0,2,3])**2).item() +def avg_sq_ch_mean(model, input, output): + """ calculate average channel square mean of output activations + """ + return torch.mean(output.mean(axis=[0, 2, 3]) ** 2).item() -def avg_ch_var(model, input, output): - "calculate average channel variance of output activations" - return torch.mean(output.var(axis=[0,2,3])).item()\ +def avg_ch_var(model, input, output): + """calculate average channel variance of output activations + """ + return torch.mean(output.var(axis=[0, 2, 3])).item() -def avg_ch_var_residual(model, input, output): - "calculate average channel variance of output activations" - return torch.mean(output.var(axis=[0,2,3])).item() +def avg_ch_var_residual(model, input, output): + """calculate average channel variance of output activations + """ + return torch.mean(output.var(axis=[0, 2, 3])).item() class ActivationStatsHook: @@ -62,15 +67,16 @@ class ActivationStatsHook: raise ValueError("Please provide `hook_fns` for each `hook_fn_locs`, \ their lengths are different.") self.stats = dict((hook_fn.__name__, []) for hook_fn in hook_fns) - for hook_fn_loc, hook_fn in zip(hook_fn_locs, hook_fns): + for hook_fn_loc, hook_fn in zip(hook_fn_locs, hook_fns): self.register_hook(hook_fn_loc, hook_fn) def _create_hook(self, hook_fn): def append_activation_stats(module, input, output): out = hook_fn(module, input, output) self.stats[hook_fn.__name__].append(out) + return append_activation_stats - + def register_hook(self, hook_fn_loc, hook_fn): for name, module in self.model.named_modules(): if not fnmatch.fnmatch(name, hook_fn_loc): @@ -78,17 +84,18 @@ class ActivationStatsHook: module.register_forward_hook(self._create_hook(hook_fn)) -def extract_spp_stats(model, - hook_fn_locs, - hook_fns, - input_shape=[8, 3, 224, 224]): +def extract_spp_stats( + model, + hook_fn_locs, + hook_fns, + input_shape=[8, 3, 224, 224]): """Extract average square channel mean and variance of activations during forward pass to plot Signal Propogation Plots (SPP). Paper: https://arxiv.org/abs/2101.08692 Example Usage: https://gist.github.com/amaarora/6e56942fcb46e67ba203f3009b30d950 - """ + """ x = torch.normal(0., 1., input_shape) hook = ActivationStatsHook(model, hook_fn_locs=hook_fn_locs, hook_fns=hook_fns) _ = model(x) @@ -186,7 +193,7 @@ def _freeze_unfreeze(root_module, submodules=[], include_bn_running_stats=True, named_modules = submodules submodules = [root_module.get_submodule(m) for m in submodules] - if not(len(submodules)): + if not (len(submodules)): named_modules, submodules = list(zip(*root_module.named_children())) for n, m in zip(named_modules, submodules): @@ -201,13 +208,14 @@ def _freeze_unfreeze(root_module, submodules=[], include_bn_running_stats=True, module.get_submodule(split[0]).add_module(split[1], submodule) else: module.add_module(name, submodule) + # Freeze batch norm if mode == 'freeze': res = freeze_batch_norm_2d(m) # It's possible that `m` is a type of BatchNorm in itself, in which case `unfreeze_batch_norm_2d` won't # convert it in place, but will return the converted result. In this case `res` holds the converted # result and we may try to re-assign the named module - if isinstance(m, (torch.nn.modules.batchnorm.BatchNorm2d, torch.nn.modules.batchnorm.SyncBatchNorm)): + if isinstance(m, (torch.nn.modules.batchnorm.BatchNorm2d, torch.nn.modules.batchnorm.SyncBatchNorm)): _add_submodule(root_module, n, res) # Unfreeze batch norm else: