eric-haibin-lin commented on a change in pull request #11502: [MXNET-614]
Adding Synchronized Batch Normalization
URL: https://github.com/apache/incubator-mxnet/pull/11502#discussion_r200206864
##########
File path: python/mxnet/gluon/contrib/nn/basic_layers.py
##########
@@ -151,3 +153,77 @@ def __repr__(self):
s = '{block_name}({input_dim} -> {output_dim}, {dtype})'
return s.format(block_name=self.__class__.__name__,
**self._kwargs)
+
+class SyncBatchNorm(BatchNorm):
+ """Cross-GPU Synchronized Batch normalization (SyncBN)
+ Standard BN [1]_ implementation only normalize the data within each device.
+ SyncBN normalizes the input within the whole mini-batch.
+ We follow the sync-onece implmentation described in the paper [2]_ .
+ Parameters
+ ----------
+ in_channels : int, default 0
+ Number of channels (feature maps) in input data. If not specified,
+ initialization will be deferred to the first time `forward` is called
+ and `in_channels` will be inferred from the shape of input data.
+ num_devices : int, default number of visible GPUs
+ momentum: float, default 0.9
+ Momentum for the moving average.
+ epsilon: float, default 1e-5
+ Small float added to variance to avoid dividing by zero.
+ center: bool, default True
+ If True, add offset of `beta` to normalized tensor.
+ If False, `beta` is ignored.
+ scale: bool, default True
+ If True, multiply by `gamma`. If False, `gamma` is not used.
+ When the next layer is linear (also e.g. `nn.relu`),
+ this can be disabled since the scaling
+ will be done by the next layer.
+ use_global_stats: bool, default False
+ If True, use global moving statistics instead of local batch-norm.
This will force
+ change batch-norm into a scale shift operator.
+ If False, use local batch-norm.
+ beta_initializer: str or `Initializer`, default 'zeros'
+ Initializer for the beta weight.
+ gamma_initializer: str or `Initializer`, default 'ones'
+ Initializer for the gamma weight.
+ moving_mean_initializer: str or `Initializer`, default 'zeros'
+ Initializer for the moving mean.
+ moving_variance_initializer: str or `Initializer`, default 'ones'
+ Initializer for the moving variance.
+
+
+ Inputs:
+ - **data**: input tensor with arbitrary shape.
+ Outputs:
+ - **out**: output tensor with the same shape as `data`.
+
+ Reference:
+ .. [1] Ioffe, Sergey, and Christian Szegedy. "Batch normalization:
Accelerating
+ deep network training by reducing internal covariate shift." *ICML
2015*
+ .. [2] Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang,
Xiaogang Wang,
+ Ambrish Tyagi, and Amit Agrawal. "Context Encoding for Semantic
Segmentation." *CVPR 2018*
+ """
+ def __init__(self, in_channels=0, num_devices=None, momentum=0.9,
epsilon=1e-5,
+ center=True, scale=True, use_global_stats=False,
beta_initializer='zeros',
+ gamma_initializer='ones', running_mean_initializer='zeros',
+ running_variance_initializer='ones', **kwargs):
+ super(SyncBatchNorm, self).__init__(1, momentum, epsilon, center,
scale, use_global_stats,
+ beta_initializer,
gamma_initializer,
+ running_mean_initializer,
running_variance_initializer,
+ in_channels, **kwargs)
+ num_devices = self._get_num_devices() if num_devices is None else
num_devices
+ self._kwargs = {'eps': epsilon, 'momentum': momentum,
+ 'fix_gamma': not scale, 'use_global_stats':
use_global_stats,
+ 'ndev': num_devices, 'key': self.prefix}
+
+ def _get_num_devices(self):
+ warnings.warn("Caution using SyncBatchNorm: "
+ "if not using all the GPUs, please mannually set
num_devices",
+ UserWarning)
+ num_devices = len(test_utils.list_gpus())
+ num_devices = num_devices if num_devices > 0 else 1
+ return num_devices
+
+ def hybrid_forward(self, F, x, gamma, beta, running_mean, running_var):
+ return F.SyncBatchNorm(x, gamma, beta, running_mean, running_var,
Review comment:
F.contrib.SyncBatchNorm
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
