| |
|
| | from copy import deepcopy
|
| |
|
| | import torch
|
| | import torch.nn as nn
|
| |
|
| | from mmaction.registry import MODELS
|
| |
|
| |
|
| | @MODELS.register_module()
|
| | class SubBatchNorm3D(nn.Module):
|
| | """Sub BatchNorm3d splits the batch dimension into N splits, and run BN on
|
| | each of them separately (so that the stats are computed on each subset of
|
| | examples (1/N of batch) independently). During evaluation, it aggregates
|
| | the stats from all splits into one BN.
|
| |
|
| | Args:
|
| | num_features (int): Dimensions of BatchNorm.
|
| | """
|
| |
|
| | def __init__(self, num_features, **cfg):
|
| | super(SubBatchNorm3D, self).__init__()
|
| |
|
| | self.num_features = num_features
|
| | self.cfg_ = deepcopy(cfg)
|
| | self.num_splits = self.cfg_.pop('num_splits', 1)
|
| | self.num_features_split = self.num_features * self.num_splits
|
| |
|
| | self.cfg_['affine'] = False
|
| | self.bn = nn.BatchNorm3d(num_features, **self.cfg_)
|
| | self.split_bn = nn.BatchNorm3d(self.num_features_split, **self.cfg_)
|
| | self.init_weights(cfg)
|
| |
|
| | def init_weights(self, cfg):
|
| | """Initialize weights."""
|
| | if cfg.get('affine', True):
|
| | self.weight = torch.nn.Parameter(torch.ones(self.num_features))
|
| | self.bias = torch.nn.Parameter(torch.zeros(self.num_features))
|
| | self.affine = True
|
| | else:
|
| | self.affine = False
|
| |
|
| | def _get_aggregated_mean_std(self, means, stds, n):
|
| | """Calculate aggregated mean and std."""
|
| | mean = means.view(n, -1).sum(0) / n
|
| | std = stds.view(n, -1).sum(0) / n + (
|
| | (means.view(n, -1) - mean)**2).view(n, -1).sum(0) / n
|
| | return mean.detach(), std.detach()
|
| |
|
| | def aggregate_stats(self):
|
| | """Synchronize running_mean, and running_var to self.bn.
|
| |
|
| | Call this before eval, then call model.eval(); When eval, forward
|
| | function will call self.bn instead of self.split_bn, During this time
|
| | the running_mean, and running_var of self.bn has been obtained from
|
| | self.split_bn.
|
| | """
|
| | if self.split_bn.track_running_stats:
|
| | aggre_func = self._get_aggregated_mean_std
|
| | self.bn.running_mean.data, self.bn.running_var.data = aggre_func(
|
| | self.split_bn.running_mean, self.split_bn.running_var,
|
| | self.num_splits)
|
| | self.bn.num_batches_tracked = self.split_bn.num_batches_tracked.detach(
|
| | )
|
| |
|
| | def forward(self, x):
|
| | """Defines the computation performed at every call."""
|
| | if self.training:
|
| | n, c, t, h, w = x.shape
|
| | assert n % self.num_splits == 0
|
| | x = x.view(n // self.num_splits, c * self.num_splits, t, h, w)
|
| | x = self.split_bn(x)
|
| | x = x.view(n, c, t, h, w)
|
| | else:
|
| | x = self.bn(x)
|
| | if self.affine:
|
| | x = x * self.weight.view(-1, 1, 1, 1)
|
| | x = x + self.bias.view(-1, 1, 1, 1)
|
| | return x
|
| |
|
