Post merge cleanup, restore previous unwrap fn

timm/utils/model.py

@@ -7,14 +7,16 @@ 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):
@@ -22,18 +24,21 @@ def get_state_dict(model, unwrap_fn=unwrap_model):
 
 
 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()
+    """ 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()\
+    """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()
+    """calculate average channel variance of output activations
+    """
+    return torch.mean(output.var(axis=[0, 2, 3])).item()
 
 
 class ActivationStatsHook:
@@ -69,6 +74,7 @@ class ActivationStatsHook:
         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):
@@ -78,7 +84,8 @@ class ActivationStatsHook:
             module.register_forward_hook(self._create_hook(hook_fn))
 
 
-def extract_spp_stats(model, 
+def extract_spp_stats(
+        model,
         hook_fn_locs,
         hook_fns,
         input_shape=[8, 3, 224, 224]):
@@ -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,6 +208,7 @@ 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)