fix the implementation of WS (fix #3)

ppwwyyxx · Jun 28, 2019 · 58f4903 · 58f4903
1 parent 9726d7f
commit 58f4903
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ImageNet-ResNet-TensorFlow/imagenet-resnet-gn.py b/ImageNet-ResNet-TensorFlow/imagenet-resnet-gn.py
@@ -41,8 +41,8 @@ def weight_standardization(v):
                 return v
             if (not v.name.endswith('/W:0')) or v.shape.ndims != 4:
                 return v
-            mean, std = tf.nn.moments(v, [0, 1, 2], keep_dims=True)
-            v = (v - mean) / (std + 1e-5)
+            mean, var = tf.nn.moments(v, [0, 1, 2], keep_dims=True)
+            v = (v - mean) / (tf.sqrt(var)+ 1e-5)
             return v
 
         num_blocks = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3]}[self.depth]