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