oneflow.nn.SyncBatchNorm¶
-
class
oneflow.nn.
SyncBatchNorm
(num_features: int, eps: float = 1e-05, momentum: float = 0.1, affine: bool = True, track_running_stats: bool = True)¶ Applies Batch Normalization over a N-Dimensional input (a mini-batch of [N-2]D inputs with additional channel dimension) as described in the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift .
\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]The mean and standard-deviation are calculated per-dimension over all mini-batches of the same process groups. \(\gamma\) and \(\beta\) are learnable parameter vectors of size C (where C is the input size). By default, the elements of \(\gamma\) are sampled from \(\mathcal{U}(0, 1)\) and the elements of \(\beta\) are set to 0. The standard-deviation is calculated via the biased estimator, equivalent to oneflow.var(input, unbiased=False).
Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default
momentum
of 0.1.If
track_running_stats
is set toFalse
, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well.Note
This
momentum
argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is \(\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t\), where \(\hat{x}\) is the estimated statistic and \(x_t\) is the new observed value.Because the Batch Normalization is done for each channel in the
C
dimension, computing statistics on(N, +)
slices, it’s common terminology to call this Volumetric Batch Normalization or Spatio-temporal Batch Normalization.Currently
SyncBatchNorm
only supportsDistributedDataParallel
(DDP) with single GPU per process. Useoneflow.nn.SyncBatchNorm.convert_sync_batchnorm()
to convertBatchNorm*D
layer toSyncBatchNorm
before wrapping Network with DDP.- Parameters
num_features – \(C\) from an expected input of size \((N, C, +)\)
eps – a value added to the denominator for numerical stability. Default:
1e-5
momentum – the value used for the running_mean and running_var computation. Can be set to
None
for cumulative moving average (i.e. simple average). Default: 0.1affine – a boolean value that when set to
True
, this module has learnable affine parameters. Default:True
track_running_stats – a boolean value that when set to
True
, this module tracks the running mean and variance, and when set toFalse
, this module does not track such statistics, and initializes statistics buffersrunning_mean
andrunning_var
asNone
. When these buffers areNone
, this module always uses batch statistics. in both training and eval modes. Default:True
- Shape:
Input: \((N, C, +)\)
Output: \((N, C, +)\) (same shape as input)
Note
Synchronization of batchnorm statistics occurs only while training, i.e. synchronization is disabled when
model.eval()
is set or ifself.training
is otherwiseFalse
.Examples:
>>> import oneflow as flow >>> bn = flow.nn.BatchNorm2d(100) >>> sync_bn = flow.nn.SyncBatchNorm.convert_sync_batchnorm(bn).cuda() >>> input = flow.randn(20, 100, 35, 45, device="cuda") >>> output = sync_bn(input)