sync_batch_norm

dragon.vm.torch.nn.functional.sync_batch_norm(
  input,
  running_mean,
  running_var,
  weight,
  bias,
  training=False,
  momentum=0.1,
  eps=1e-05,
  process_group=None
)[source]

Apply the sync batch normalization over input. [Ioffe & Szegedy, 2015].

The normalization is defined as:

\[y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta \]

The moving average of stats are calculated as:

\[x_{moving} \leftarrow (1 - momentum) * x_{moving} + momentum * x_{stat} \]

Additionally, you can specify process_group to perform synchronization.

If not, value returning from dragon.distributed.get_group(...) will be used.

Parameters:
  • input (dragon.vm.torch.Tensor) – The input tensor.
  • running_mean (dragon.vm.torch.Tensor) – The running mean.
  • running_var (dragon.vm.torch.Tensor) – The running var.
  • weight (dragon.vm.torch.Tensor) – The weight tensor.
  • bias (dragon.vm.torch.Tensor) – The bias tensor.
  • training (bool, optional, default=False) – The flag to determine the stats.
  • momentum (float, optional, default=0.1) – The momentum to the moving average.
  • eps (float, optional, default=1e-5) – The epsilon value.
  • process_group (dragon.ProcessGroup, optional) – The group for communication.
Returns:

dragon.vm.torch.Tensor – The output tensor.