You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
154 lines
6.2 KiB
154 lines
6.2 KiB
""" Model / state_dict utils
|
|
|
|
Hacked together by / Copyright 2020 Ross Wightman
|
|
"""
|
|
from logging import root
|
|
from typing import Sequence
|
|
import re
|
|
import warnings
|
|
|
|
import torch
|
|
import fnmatch
|
|
|
|
from torch.nn.modules import module
|
|
|
|
from .model_ema import ModelEma
|
|
from timm.models.layers.norm import FrozenBatchNorm2d
|
|
|
|
|
|
def unwrap_model(model):
|
|
if isinstance(model, ModelEma):
|
|
return unwrap_model(model.ema)
|
|
else:
|
|
return model.module if hasattr(model, 'module') else 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()\
|
|
|
|
|
|
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 matches modules using unix pattern
|
|
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
|
|
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`.
|
|
|
|
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):
|
|
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)
|
|
|
|
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):
|
|
continue
|
|
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]):
|
|
"""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)
|
|
return hook.stats
|
|
|
|
|
|
def freeze(modules, root_module=None, include_bn_running_stats=True, mode=True):
|
|
"""
|
|
Freeze or unfreeze parameters of the specified modules and those of all their hierarchical descendants. This is
|
|
done in place.
|
|
Args:
|
|
modules (nn.Module or list[nn.Module] or str or list[str]): List of modules for which the parameters will be
|
|
(un)frozen. If a string or strings are provided these will be interpreted according to the named modules
|
|
of the provided ``root_module``.
|
|
root_module (nn.Module, optional): Root module relative to which named modules (accessible via
|
|
``root_module.named_modules()``) are referenced. Must be provided if the `modules` argument is specified
|
|
with a string or strings. Defaults to `None`.
|
|
include_bn_running_stats (bool): Whether to also (un)freeze the running statistics of batch norm layers.
|
|
Defaults to `True`.
|
|
mode (bool): Whether to freeze (`True`) or unfreeze (`False`). Defaults to `True`.
|
|
|
|
TODO before finalizing PR: Implement unfreezing of batch norm
|
|
"""
|
|
|
|
if not isinstance(modules, Sequence):
|
|
modules = [modules]
|
|
|
|
if isinstance(modules[0], str):
|
|
assert root_module is not None, \
|
|
"When providing strings for the `modules` argument, a `root_module` must be provided"
|
|
module_names = modules
|
|
modules = [root_module.get_submodule(m) for m in module_names]
|
|
|
|
for n, m in zip(module_names, modules):
|
|
for p in m.parameters():
|
|
p.requires_grad = (not mode)
|
|
if include_bn_running_stats:
|
|
res = FrozenBatchNorm2d.convert_frozen_batchnorm(m)
|
|
# It's possible that `m` is a type of BatchNorm in itself, in which case
|
|
# `FrozenBatchNorm2d.convert_frozen_batchnorm` 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 module_names is not None and root_module is not None:
|
|
root_module.add_module(n, res)
|
|
else:
|
|
raise RuntimeError(
|
|
"Could not freeze batch norm statistics due to a technical limitation. Hint: Try calling "
|
|
"`freeze` with a list of module names while providing a `root_module` argument.")
|
|
|
|
|
|
def unfreeze(modules, root_module=None, include_bn_running_stats=True):
|
|
"""
|
|
Idiomatic convenience function to call `freeze` with `mode == False`. See docstring of `freeze` for further
|
|
information.
|
|
"""
|
|
freeze(modules, root_module=root_module, include_bn_running_stats=include_bn_running_stats, mode=False)
|