# encoding: utf-8
|
"""
|
@author: liaoxingyu
|
@contact: sherlockliao01@gmail.com
|
"""
|
|
import logging
|
|
import torch
|
import torch.nn.functional as F
|
from torch import nn
|
|
__all__ = [
|
"BatchNorm",
|
"IBN",
|
"GhostBatchNorm",
|
"FrozenBatchNorm",
|
"SyncBatchNorm",
|
"get_norm",
|
]
|
|
|
class BatchNorm(nn.BatchNorm2d):
|
def __init__(self, num_features, eps=1e-05, momentum=0.1, weight_freeze=False, bias_freeze=False, weight_init=1.0,
|
bias_init=0.0, **kwargs):
|
super().__init__(num_features, eps=eps, momentum=momentum)
|
if weight_init is not None: nn.init.constant_(self.weight, weight_init)
|
if bias_init is not None: nn.init.constant_(self.bias, bias_init)
|
self.weight.requires_grad_(not weight_freeze)
|
self.bias.requires_grad_(not bias_freeze)
|
|
|
class SyncBatchNorm(nn.SyncBatchNorm):
|
def __init__(self, num_features, eps=1e-05, momentum=0.1, weight_freeze=False, bias_freeze=False, weight_init=1.0,
|
bias_init=0.0):
|
super().__init__(num_features, eps=eps, momentum=momentum)
|
if weight_init is not None: nn.init.constant_(self.weight, weight_init)
|
if bias_init is not None: nn.init.constant_(self.bias, bias_init)
|
self.weight.requires_grad_(not weight_freeze)
|
self.bias.requires_grad_(not bias_freeze)
|
|
|
class IBN(nn.Module):
|
def __init__(self, planes, bn_norm, **kwargs):
|
super(IBN, self).__init__()
|
half1 = int(planes / 2)
|
self.half = half1
|
half2 = planes - half1
|
self.IN = nn.InstanceNorm2d(half1, affine=True)
|
self.BN = get_norm(bn_norm, half2, **kwargs)
|
|
def forward(self, x):
|
split = torch.split(x, self.half, 1)
|
out1 = self.IN(split[0].contiguous())
|
out2 = self.BN(split[1].contiguous())
|
out = torch.cat((out1, out2), 1)
|
return out
|
|
|
class GhostBatchNorm(BatchNorm):
|
def __init__(self, num_features, num_splits=1, **kwargs):
|
super().__init__(num_features, **kwargs)
|
self.num_splits = num_splits
|
self.register_buffer('running_mean', torch.zeros(num_features))
|
self.register_buffer('running_var', torch.ones(num_features))
|
|
def forward(self, input):
|
N, C, H, W = input.shape
|
if self.training or not self.track_running_stats:
|
self.running_mean = self.running_mean.repeat(self.num_splits)
|
self.running_var = self.running_var.repeat(self.num_splits)
|
outputs = F.batch_norm(
|
input.view(-1, C * self.num_splits, H, W), self.running_mean, self.running_var,
|
self.weight.repeat(self.num_splits), self.bias.repeat(self.num_splits),
|
True, self.momentum, self.eps).view(N, C, H, W)
|
self.running_mean = torch.mean(self.running_mean.view(self.num_splits, self.num_features), dim=0)
|
self.running_var = torch.mean(self.running_var.view(self.num_splits, self.num_features), dim=0)
|
return outputs
|
else:
|
return F.batch_norm(
|
input, self.running_mean, self.running_var,
|
self.weight, self.bias, False, self.momentum, self.eps)
|
|
|
class FrozenBatchNorm(BatchNorm):
|
"""
|
BatchNorm2d where the batch statistics and the affine parameters are fixed.
|
It contains non-trainable buffers called
|
"weight" and "bias", "running_mean", "running_var",
|
initialized to perform identity transformation.
|
The pre-trained backbone models from Caffe2 only contain "weight" and "bias",
|
which are computed from the original four parameters of BN.
|
The affine transform `x * weight + bias` will perform the equivalent
|
computation of `(x - running_mean) / sqrt(running_var) * weight + bias`.
|
When loading a backbone model from Caffe2, "running_mean" and "running_var"
|
will be left unchanged as identity transformation.
|
Other pre-trained backbone models may contain all 4 parameters.
|
The forward is implemented by `F.batch_norm(..., training=False)`.
|
"""
|
|
_version = 3
|
|
def __init__(self, num_features, eps=1e-5, **kwargs):
|
super().__init__(num_features, weight_freeze=True, bias_freeze=True, **kwargs)
|
self.num_features = num_features
|
self.eps = eps
|
|
def forward(self, x):
|
if x.requires_grad:
|
# When gradients are needed, F.batch_norm will use extra memory
|
# because its backward op computes gradients for weight/bias as well.
|
scale = self.weight * (self.running_var + self.eps).rsqrt()
|
bias = self.bias - self.running_mean * scale
|
scale = scale.reshape(1, -1, 1, 1)
|
bias = bias.reshape(1, -1, 1, 1)
|
return x * scale + bias
|
else:
|
# When gradients are not needed, F.batch_norm is a single fused op
|
# and provide more optimization opportunities.
|
return F.batch_norm(
|
x,
|
self.running_mean,
|
self.running_var,
|
self.weight,
|
self.bias,
|
training=False,
|
eps=self.eps,
|
)
|
|
def _load_from_state_dict(
|
self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
|
):
|
version = local_metadata.get("version", None)
|
|
if version is None or version < 2:
|
# No running_mean/var in early versions
|
# This will silent the warnings
|
if prefix + "running_mean" not in state_dict:
|
state_dict[prefix + "running_mean"] = torch.zeros_like(self.running_mean)
|
if prefix + "running_var" not in state_dict:
|
state_dict[prefix + "running_var"] = torch.ones_like(self.running_var)
|
|
if version is not None and version < 3:
|
logger = logging.getLogger(__name__)
|
logger.info("FrozenBatchNorm {} is upgraded to version 3.".format(prefix.rstrip(".")))
|
# In version < 3, running_var are used without +eps.
|
state_dict[prefix + "running_var"] -= self.eps
|
|
super()._load_from_state_dict(
|
state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
|
)
|
|
def __repr__(self):
|
return "FrozenBatchNorm2d(num_features={}, eps={})".format(self.num_features, self.eps)
|
|
@classmethod
|
def convert_frozen_batchnorm(cls, module):
|
"""
|
Convert BatchNorm/SyncBatchNorm in module into FrozenBatchNorm.
|
Args:
|
module (torch.nn.Module):
|
Returns:
|
If module is BatchNorm/SyncBatchNorm, returns a new module.
|
Otherwise, in-place convert module and return it.
|
Similar to convert_sync_batchnorm in
|
https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/batchnorm.py
|
"""
|
bn_module = nn.modules.batchnorm
|
bn_module = (bn_module.BatchNorm2d, bn_module.SyncBatchNorm)
|
res = module
|
if isinstance(module, bn_module):
|
res = cls(module.num_features)
|
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 name, child in module.named_children():
|
new_child = cls.convert_frozen_batchnorm(child)
|
if new_child is not child:
|
res.add_module(name, new_child)
|
return res
|
|
|
def get_norm(norm, out_channels, **kwargs):
|
"""
|
Args:
|
norm (str or callable): either one of BN, GhostBN, FrozenBN, GN or SyncBN;
|
or a callable that thakes a channel number and returns
|
the normalization layer as a nn.Module
|
out_channels: number of channels for normalization layer
|
|
Returns:
|
nn.Module or None: the normalization layer
|
"""
|
if isinstance(norm, str):
|
if len(norm) == 0:
|
return None
|
norm = {
|
"BN": BatchNorm,
|
"GhostBN": GhostBatchNorm,
|
"FrozenBN": FrozenBatchNorm,
|
"GN": lambda channels, **args: nn.GroupNorm(32, channels),
|
"syncBN": SyncBatchNorm,
|
}[norm]
|
return norm(out_channels, **kwargs)
|
