SyncBatchNorm#

classtorch.nn.SyncBatchNorm(num_features,eps=1e-05,momentum=0.1,affine=True,track_running_stats=True,process_group=None,device=None,dtype=None)[source]#

Applies Batch Normalization over a N-Dimensional input.

The N-D input is a mini-batch of [N-2]D inputs with additional channel dimension) as described in the paperBatch Normalization: Accelerating Deep Network Training by ReducingInternal Covariate Shift .

$$y=\frac{x-\mathrm{E}[x]}{\sqrt{\mathrm{V}\mathrm{a}\mathrm{r}[x]+ϵ}}\ast \gamma +\beta$$y=Var[x]+ϵ​x−E[x]​∗γ+β

The mean and standard-deviation are calculated per-dimension over allmini-batches of the same process groups.$\gamma$γ and$\beta$βare learnable parameter vectors of sizeC (whereC is the input size).By default, the elements of$\gamma$γ are sampled from$\mathcal{U}(0,1)$U(0,1) and the elements of$\beta$β are set to 0.The standard-deviation is calculated via the biased estimator, equivalent totorch.var(input, unbiased=False).

Also by default, during training this layer keeps running estimates of itscomputed mean and variance, which are then used for normalization duringevaluation. The running estimates are kept with a defaultmomentumof 0.1.

Iftrack_running_stats is set toFalse, this layer then does notkeep running estimates, and batch statistics are instead used duringevaluation time as well.

Note

Thismomentum argument is different from one used in optimizerclasses and the conventional notion of momentum. Mathematically, theupdate rule for running statistics here is${\widehat{x}}_{\text{new}}=(1-\text{momentum})\times \widehat{x}+\text{momentum}\times x_t$x^new​=(1−momentum)×x^+momentum×xt​,where$\widehat{x}$x^ is the estimated statistic and$x_t$xt​ is thenew observed value.

Because the Batch Normalization is done for each channel in theC dimension, computingstatistics on(N,+) slices, it’s common terminology to call this Volumetric BatchNormalization or Spatio-temporal Batch Normalization.

CurrentlySyncBatchNorm only supportsDistributedDataParallel (DDP) with single GPU per process. Usetorch.nn.SyncBatchNorm.convert_sync_batchnorm() to convertBatchNorm*D layer toSyncBatchNorm before wrappingNetwork with DDP.

Parameters

• num_features (int) –$C$C from an expected input of size$(N,C,+)$(N,C,+)

• eps (float) – a value added to the denominator for numerical stability.Default:1e-5

• momentum (Optional[float]) – the value used for the running_mean and running_varcomputation. Can be set toNone for cumulative moving average(i.e. simple average). Default: 0.1

• affine (bool) – a boolean value that when set toTrue, this module haslearnable affine parameters. Default:True

• track_running_stats (bool) – a boolean value that when set toTrue, thismodule tracks the running mean and variance, and when set toFalse,this module does not track such statistics, and initializes statisticsbuffersrunning_mean andrunning_var asNone.When these buffers areNone, this module always uses batch statistics.in both training and eval modes. Default:True

• process_group (Optional[Any]) – synchronization of stats happen within each process groupindividually. Default behavior is synchronization across the wholeworld

Shape:

• Input:$(N,C,+)$(N,C,+)

• Output:$(N,C,+)$(N,C,+) (same shape as input)

Note

Synchronization of batchnorm statistics occurs only while training, i.e.synchronization is disabled whenmodel.eval() is set or ifself.training is otherwiseFalse.

Examples:

>>># With Learnable Parameters>>>m=nn.SyncBatchNorm(100)>>># creating process group (optional)>>># ranks is a list of int identifying rank ids.>>>ranks=list(range(8))>>>r1,r2=ranks[:4],ranks[4:]>>># Note: every rank calls into new_group for every>>># process group created, even if that rank is not>>># part of the group.>>>process_groups=[torch.distributed.new_group(pids)forpidsin[r1,r2]]>>>process_group=process_groups[0ifdist.get_rank()<=3else1]>>># Without Learnable Parameters>>>m=nn.BatchNorm3d(100,affine=False,process_group=process_group)>>>input=torch.randn(20,100,35,45,10)>>>output=m(input)>>># network is nn.BatchNorm layer>>>sync_bn_network=nn.SyncBatchNorm.convert_sync_batchnorm(network,process_group)>>># only single gpu per process is currently supported>>>ddp_sync_bn_network=torch.nn.parallel.DistributedDataParallel(>>>sync_bn_network,>>>device_ids=[args.local_rank],>>>output_device=args.local_rank)

classmethodconvert_sync_batchnorm(module,process_group=None)[source]#

Converts allBatchNorm*D layers in the model totorch.nn.SyncBatchNorm layers.

Parameters

• module (nn.Module) – module containing one or moreBatchNorm*D layers

• process_group (optional) – process group to scope synchronization,default is the whole world

Returns

The originalmodule with the convertedtorch.nn.SyncBatchNormlayers. If the originalmodule is aBatchNorm*D layer,a newtorch.nn.SyncBatchNorm layer object will be returnedinstead.

Example:

>>># Network with nn.BatchNorm layer>>>module=torch.nn.Sequential(>>>torch.nn.Linear(20,100),>>>torch.nn.BatchNorm1d(100),>>>).cuda()>>># creating process group (optional)>>># ranks is a list of int identifying rank ids.>>>ranks=list(range(8))>>>r1,r2=ranks[:4],ranks[4:]>>># Note: every rank calls into new_group for every>>># process group created, even if that rank is not>>># part of the group.>>>process_groups=[torch.distributed.new_group(pids)forpidsin[r1,r2]]>>>process_group=process_groups[0ifdist.get_rank()<=3else1]>>>sync_bn_module=torch.nn.SyncBatchNorm.convert_sync_batchnorm(module,process_group)

forward(input)[source]#

Runs the forward pass.

Return type

Tensor