Spaces:
Running
Running
from itertools import repeat | |
import collections.abc | |
from torch import nn as nn | |
from torchvision.ops.misc import FrozenBatchNorm2d | |
def freeze_batch_norm_2d(module, module_match={}, name=''): | |
""" | |
Converts all `BatchNorm2d` and `SyncBatchNorm` layers of provided module into `FrozenBatchNorm2d`. If `module` is | |
itself an instance of either `BatchNorm2d` or `SyncBatchNorm`, it is converted into `FrozenBatchNorm2d` and | |
returned. Otherwise, the module is walked recursively and submodules are converted in place. | |
Args: | |
module (torch.nn.Module): Any PyTorch module. | |
module_match (dict): Dictionary of full module names to freeze (all if empty) | |
name (str): Full module name (prefix) | |
Returns: | |
torch.nn.Module: Resulting module | |
Inspired by https://github.com/pytorch/pytorch/blob/a5895f85be0f10212791145bfedc0261d364f103/torch/nn/modules/batchnorm.py#L762 | |
""" | |
res = module | |
is_match = True | |
if module_match: | |
is_match = name in module_match | |
if is_match and isinstance(module, (nn.modules.batchnorm.BatchNorm2d, nn.modules.batchnorm.SyncBatchNorm)): | |
res = FrozenBatchNorm2d(module.num_features) | |
res.num_features = module.num_features | |
res.affine = module.affine | |
if module.affine: | |
res.weight.data = module.weight.data.clone().detach() | |
res.bias.data = module.bias.data.clone().detach() | |
res.running_mean.data = module.running_mean.data | |
res.running_var.data = module.running_var.data | |
res.eps = module.eps | |
else: | |
for child_name, child in module.named_children(): | |
full_child_name = '.'.join([name, child_name]) if name else child_name | |
new_child = freeze_batch_norm_2d(child, module_match, full_child_name) | |
if new_child is not child: | |
res.add_module(child_name, new_child) | |
return res | |
# From PyTorch internals | |
def _ntuple(n): | |
def parse(x): | |
if isinstance(x, collections.abc.Iterable): | |
return x | |
return tuple(repeat(x, n)) | |
return parse | |
to_1tuple = _ntuple(1) | |
to_2tuple = _ntuple(2) | |
to_3tuple = _ntuple(3) | |
to_4tuple = _ntuple(4) | |
to_ntuple = lambda n, x: _ntuple(n)(x) | |