add sparse SyncBatchNorm

PaddlePaddle · Jun 14, 2022 · f8c44d2 · f8c44d2
1 parent 084f2a9
commit f8c44d2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/python/paddle/incubate/sparse/nn/__init__.py b/python/paddle/incubate/sparse/nn/__init__.py
@@ -15,14 +15,15 @@
 from . import functional
 
 from .layer.activation import ReLU
-from .layer.norm import BatchNorm
+from .layer.norm import BatchNorm, SyncBatchNorm
 from .layer.conv import Conv3D
 from .layer.conv import SubmConv3D
 from .layer.pooling import MaxPool3D
 
 __all__ = [
     'ReLU',
     'BatchNorm',
+    'SyncBatchNorm',
     'Conv3D',
     'SubmConv3D',
     'MaxPool3D',

diff --git a/python/paddle/incubate/sparse/nn/layer/norm.py b/python/paddle/incubate/sparse/nn/layer/norm.py
@@ -27,6 +27,8 @@
 
 import paddle
 import warnings
+from paddle.nn.layer.norm import _BatchNormBase
+from paddle.framework import no_grad
 
 
 class BatchNorm(paddle.nn.BatchNorm1D):
@@ -171,7 +173,7 @@ def __init__(self,
                  name=None):
         super(SyncBatchNorm,
               self).__init__(num_features, momentum, epsilon, weight_attr,
-                             bias_attr, data_format, None, name)
+                             bias_attr, data_format, name)
 
     def forward(self, x):
         out = super(SyncBatchNorm, self).forward(x.values())