|
|
|
|
|
|
|
"""BatchNorm (BN) utility functions and custom batch-size BN implementations""" |
|
|
|
from functools import partial |
|
|
|
import torch |
|
import torch.nn as nn |
|
|
|
from pytorchvideo.layers.batch_norm import NaiveSyncBatchNorm3d |
|
|
|
|
|
def get_norm(cfg): |
|
""" |
|
Args: |
|
cfg (CfgNode): model building configs, details are in the comments of |
|
the config file. |
|
Returns: |
|
nn.Module: the normalization layer. |
|
""" |
|
if cfg.BN.NORM_TYPE in {"batchnorm", "sync_batchnorm_apex"}: |
|
return nn.BatchNorm3d |
|
elif cfg.BN.NORM_TYPE == "sub_batchnorm": |
|
return partial(SubBatchNorm3d, num_splits=cfg.BN.NUM_SPLITS) |
|
elif cfg.BN.NORM_TYPE == "sync_batchnorm": |
|
return partial( |
|
NaiveSyncBatchNorm3d, |
|
num_sync_devices=cfg.BN.NUM_SYNC_DEVICES, |
|
global_sync=cfg.BN.GLOBAL_SYNC, |
|
) |
|
else: |
|
raise NotImplementedError( |
|
"Norm type {} is not supported".format(cfg.BN.NORM_TYPE) |
|
) |
|
|
|
|
|
class SubBatchNorm3d(nn.Module): |
|
""" |
|
The standard BN layer computes stats across all examples in a GPU. In some |
|
cases it is desirable to compute stats across only a subset of examples |
|
(e.g., in multigrid training https://arxiv.org/abs/1912.00998). |
|
SubBatchNorm3d 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. |
|
""" |
|
|
|
def __init__(self, num_splits, **args): |
|
""" |
|
Args: |
|
num_splits (int): number of splits. |
|
args (list): other arguments. |
|
""" |
|
super(SubBatchNorm3d, self).__init__() |
|
self.num_splits = num_splits |
|
num_features = args["num_features"] |
|
|
|
if args.get("affine", True): |
|
self.affine = True |
|
args["affine"] = False |
|
self.weight = torch.nn.Parameter(torch.ones(num_features)) |
|
self.bias = torch.nn.Parameter(torch.zeros(num_features)) |
|
else: |
|
self.affine = False |
|
self.bn = nn.BatchNorm3d(**args) |
|
args["num_features"] = num_features * num_splits |
|
self.split_bn = nn.BatchNorm3d(**args) |
|
|
|
def _get_aggregated_mean_std(self, means, stds, n): |
|
""" |
|
Calculate the aggregated mean and stds. |
|
Args: |
|
means (tensor): mean values. |
|
stds (tensor): standard deviations. |
|
n (int): number of sets of means and stds. |
|
""" |
|
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. Call this before eval. |
|
""" |
|
if self.split_bn.track_running_stats: |
|
( |
|
self.bn.running_mean.data, |
|
self.bn.running_var.data, |
|
) = self._get_aggregated_mean_std( |
|
self.split_bn.running_mean, |
|
self.split_bn.running_var, |
|
self.num_splits, |
|
) |
|
|
|
def forward(self, x): |
|
if self.training: |
|
n, c, t, h, w = x.shape |
|
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 |
|
|
