From 20626e838796136557a6b5174344a5ab5fc20973 Mon Sep 17 00:00:00 2001 From: Aman Arora Date: Sat, 27 Mar 2021 05:40:04 +1100 Subject: [PATCH 1/4] Add to extract stats for SPP --- timm/utils/model.py | 64 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/timm/utils/model.py b/timm/utils/model.py index cfd42806..856bb981 100644 --- a/timm/utils/model.py +++ b/timm/utils/model.py @@ -3,6 +3,7 @@ Hacked together by / Copyright 2020 Ross Wightman """ from .model_ema import ModelEma +import torch def unwrap_model(model): @@ -14,3 +15,66 @@ def unwrap_model(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_ch_var(model, input, output): + "calculate average channel variance of output activations" + return torch.mean(output.var(axis=[0,2,3])).item() + + +class ActivationStatsHook: + """Iterates through each of `model`'s modules and if module's class name + is present in `layer_names` then registers `hook_fns` inside that module + and stores activation stats inside `self.stats`. + + Arguments: + model (nn.Module): model from which we will extract the activation stats + layer_names (List[str]): The layer name to look for to register forward + hook. Example, `BasicBlock`, `Bottleneck` + hook_fns (List[Callable]): List of hook functions to be registered at every + module in `layer_names`. + + Inspiration from https://docs.fast.ai/callback.hook.html. + """ + + def __init__(self, model, layer_names, hook_fns=[avg_sq_ch_mean, avg_ch_var]): + self.model = model + self.layer_names = layer_names + self.hook_fns = hook_fns + self.stats = dict((hook_fn.__name__, []) for hook_fn in hook_fns) + for hook_fn in hook_fns: + self.register_hook(layer_names, 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, layer_names, hook_fn): + for layer in self.model.modules(): + layer_name = layer.__class__.__name__ + if layer_name not in layer_names: + continue + layer.register_forward_hook(self._create_hook(hook_fn)) + + +def extract_spp_stats(model, + layer_names, + hook_fns=[avg_sq_ch_mean, avg_ch_var], + 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 + """ + x = torch.normal(0., 1., input_shape) + hook = ActivationStatsHook(model, layer_names, hook_fns) + _ = model(x) + return hook.stats + \ No newline at end of file From b85be24054728179691f260a6491a1b3a03fd18c Mon Sep 17 00:00:00 2001 From: Aman Arora Date: Mon, 29 Mar 2021 09:36:31 +1100 Subject: [PATCH 2/4] update to work with fnmatch --- timm/utils/model.py | 40 ++++++++++++++++++++++------------------ 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/timm/utils/model.py b/timm/utils/model.py index 856bb981..ef2236ca 100644 --- a/timm/utils/model.py +++ b/timm/utils/model.py @@ -4,7 +4,7 @@ Hacked together by / Copyright 2020 Ross Wightman """ from .model_ema import ModelEma import torch - +import fnmatch def unwrap_model(model): if isinstance(model, ModelEma): @@ -23,32 +23,37 @@ def avg_sq_ch_mean(model, input, output): 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() class ActivationStatsHook: - """Iterates through each of `model`'s modules and if module's class name - is present in `layer_names` then registers `hook_fns` inside that module - and stores activation stats inside `self.stats`. + """Iterates through each of `model`'s modules and matches modules using unix pattern + matching based on `layer_name` and `layer_type`. If there is match, this class adds + creates a hook using `hook_fn` and adds it to the module. Arguments: model (nn.Module): model from which we will extract the activation stats - layer_names (List[str]): The layer name to look for to register forward - hook. Example, `BasicBlock`, `Bottleneck` + layer_names (str): The layer name to look for to register forward + hook. Example, 'stem', 'stages' hook_fns (List[Callable]): List of hook functions to be registered at every module in `layer_names`. Inspiration from https://docs.fast.ai/callback.hook.html. """ - def __init__(self, model, layer_names, hook_fns=[avg_sq_ch_mean, avg_ch_var]): + def __init__(self, model, hook_fn_locs, hook_fns): self.model = model - self.layer_names = layer_names + self.hook_fn_locs = hook_fn_locs self.hook_fns = hook_fns self.stats = dict((hook_fn.__name__, []) for hook_fn in hook_fns) - for hook_fn in hook_fns: - self.register_hook(layer_names, hook_fn) + 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): @@ -56,17 +61,16 @@ class ActivationStatsHook: self.stats[hook_fn.__name__].append(out) return append_activation_stats - def register_hook(self, layer_names, hook_fn): - for layer in self.model.modules(): - layer_name = layer.__class__.__name__ - if layer_name not in layer_names: + 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): continue - layer.register_forward_hook(self._create_hook(hook_fn)) + module.register_forward_hook(self._create_hook(hook_fn)) def extract_spp_stats(model, - layer_names, - hook_fns=[avg_sq_ch_mean, avg_ch_var], + 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). @@ -74,7 +78,7 @@ def extract_spp_stats(model, Paper: https://arxiv.org/abs/2101.08692 """ x = torch.normal(0., 1., input_shape) - hook = ActivationStatsHook(model, layer_names, hook_fns) + hook = ActivationStatsHook(model, hook_fn_locs=hook_fn_locs, hook_fns=hook_fns) _ = model(x) return hook.stats \ No newline at end of file From 92b1db9a79bbd9c05ad5699397a99c203e1a9a70 Mon Sep 17 00:00:00 2001 From: Aman Arora Date: Mon, 29 Mar 2021 10:04:51 +1100 Subject: [PATCH 3/4] update docstrings and add check on and --- timm/utils/model.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/timm/utils/model.py b/timm/utils/model.py index ef2236ca..4929670f 100644 --- a/timm/utils/model.py +++ b/timm/utils/model.py @@ -34,13 +34,13 @@ def avg_ch_var_residual(model, input, output): class ActivationStatsHook: """Iterates through each of `model`'s modules and matches modules using unix pattern - matching based on `layer_name` and `layer_type`. If there is match, this class adds - creates a hook using `hook_fn` and adds it to the module. + matching based on `hook_fn_locs` and registers `hook_fn` to the module if there is + a match. Arguments: model (nn.Module): model from which we will extract the activation stats - layer_names (str): The layer name to look for to register forward - hook. Example, 'stem', 'stages' + hook_fn_locs (List[str]): List of `hook_fn` locations based on Unix type string + matching with the name of model's modules. hook_fns (List[Callable]): List of hook functions to be registered at every module in `layer_names`. @@ -51,6 +51,9 @@ class ActivationStatsHook: self.model = model self.hook_fn_locs = hook_fn_locs self.hook_fns = hook_fns + if len(hook_fn_locs) != len(hook_fns): + 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): self.register_hook(hook_fn_loc, hook_fn) From 6b1806177314ef13bd93ce23be39748ae4a9ca67 Mon Sep 17 00:00:00 2001 From: Aman Arora Date: Mon, 29 Mar 2021 15:33:31 +1100 Subject: [PATCH 4/4] Add GIST to docstring for quick access --- timm/utils/model.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/timm/utils/model.py b/timm/utils/model.py index 4929670f..bd46e2f4 100644 --- a/timm/utils/model.py +++ b/timm/utils/model.py @@ -45,6 +45,9 @@ class ActivationStatsHook: module in `layer_names`. Inspiration from https://docs.fast.ai/callback.hook.html. + + Refer to https://gist.github.com/amaarora/6e56942fcb46e67ba203f3009b30d950 for an example + on how to plot Signal Propogation Plots using `ActivationStatsHook`. """ def __init__(self, model, hook_fn_locs, hook_fns): @@ -79,6 +82,8 @@ def extract_spp_stats(model, 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)