140 lines
5.2 KiB
Python
140 lines
5.2 KiB
Python
import collections.abc
|
|
from itertools import repeat
|
|
from typing import List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
from torch import nn as nn
|
|
from torch import _assert
|
|
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)
|
|
|
|
# Replaces all linear layers with linear_replacement
|
|
# TODO: add int8 support for other linear layers including attn and convnets
|
|
def replace_linear(model, linear_replacement, include_modules=['c_fc', 'c_proj'], copy_weights=True):
|
|
for name, module in model.named_children():
|
|
if len(list(module.children())) > 0:
|
|
replace_linear(module, linear_replacement, include_modules, copy_weights)
|
|
|
|
if isinstance(module, torch.nn.Linear) and name in include_modules:
|
|
old_module = model._modules[name]
|
|
model._modules[name] = linear_replacement(
|
|
module.in_features,
|
|
module.out_features,
|
|
module.bias is not None,
|
|
)
|
|
if copy_weights:
|
|
model._modules[name].weight.data.copy_(old_module.weight.data)
|
|
if model._modules[name].bias is not None:
|
|
model._modules[name].bias.data.copy_(old_module.bias)
|
|
|
|
return model
|
|
|
|
def convert_int8_model_to_inference_mode(model):
|
|
for m in model.modules():
|
|
if hasattr(m, 'prepare_for_eval'):
|
|
int8_original_dtype = m.weight.dtype
|
|
m.prepare_for_eval()
|
|
m.int8_original_dtype = int8_original_dtype
|
|
|
|
|
|
def feature_take_indices(
|
|
num_features: int,
|
|
indices: Optional[Union[int, List[int]]] = None,
|
|
as_set: bool = False,
|
|
) -> Tuple[List[int], int]:
|
|
""" Determine the absolute feature indices to 'take' from.
|
|
|
|
Note: This function can be called in forward() so must be torchscript compatible,
|
|
which requires some incomplete typing and workaround hacks.
|
|
|
|
Args:
|
|
num_features: total number of features to select from
|
|
indices: indices to select,
|
|
None -> select all
|
|
int -> select last n
|
|
list/tuple of int -> return specified (-ve indices specify from end)
|
|
as_set: return as a set
|
|
|
|
Returns:
|
|
List (or set) of absolute (from beginning) indices, Maximum index
|
|
"""
|
|
if indices is None:
|
|
indices = num_features # all features if None
|
|
|
|
if isinstance(indices, int):
|
|
# convert int -> last n indices
|
|
_assert(0 < indices <= num_features, f'last-n ({indices}) is out of range (1 to {num_features})')
|
|
take_indices = [num_features - indices + i for i in range(indices)]
|
|
else:
|
|
take_indices: List[int] = []
|
|
for i in indices:
|
|
idx = num_features + i if i < 0 else i
|
|
_assert(0 <= idx < num_features, f'feature index {idx} is out of range (0 to {num_features - 1})')
|
|
take_indices.append(idx)
|
|
|
|
if not torch.jit.is_scripting() and as_set:
|
|
return set(take_indices), max(take_indices)
|
|
|
|
return take_indices, max(take_indices)
|
|
|
|
|
|
def _out_indices_as_tuple(x: Union[int, Tuple[int, ...]]) -> Tuple[int, ...]:
|
|
if isinstance(x, int):
|
|
# if indices is an int, take last N features
|
|
return tuple(range(-x, 0))
|
|
return tuple(x)
|