mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Fix some formatting in utils/model.py

Browse Source
This commit is contained in:
Ross Wightman 2021-10-23 20:35:36 -07:00
parent 0fe4fd3f1f
commit 57992509f9
1 changed files with 26 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
timm/utils/model.py
View File

@ -2,11 +2,9 @@
Hacked together by / Copyright 2020 Ross Wightman
"""
from logging import root
from typing import Sequence
import fnmatch
import torch
import fnmatch
from torchvision.ops.misc import FrozenBatchNorm2d
from .model_ema import ModelEma
@ -24,17 +22,20 @@ 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"
""" 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"
""" calculate average channel variance of output activations
"""
return torch.mean(output.var(axis=[0, 2, 3])).item()
@ -71,6 +72,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):
@ -80,7 +82,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]):
@ -188,7 +191,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):
@ -203,6 +206,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)
@ -267,4 +271,3 @@ def unfreeze(root_module, submodules=[], include_bn_running_stats=True):
See example in docstring for `freeze`.
"""
_freeze_unfreeze(root_module, submodules, include_bn_running_stats=include_bn_running_stats, mode="unfreeze")