Merge pull request #116 from szagoruyko/bn-reset-fix

add forgotten reset for Batch Normalization

1 files changed, 11 insertions, 0 deletions

diff --git a/SpatialBatchNormalization.lua b/SpatialBatchNormalization.lua
index 093e20b..53e8f7e 100644
--- a/SpatialBatchNormalization.lua
+++ b/SpatialBatchNormalization.lua
@@ -30,6 +30,17 @@ function SpatialBatchNormalization:__init(nFeature, eps, momentum, affine)
    self.mode = 'CUDNN_BATCHNORM_SPATIAL'
 end
 
+function SpatialBatchNormalization:reset()
+   if self.weight then
+      self.weight:uniform()
+   end
+   if self.bias then
+      self.bias:zero()
+   end
+   self.running_mean:zero()
+   self.running_std:fill(1)
+end
+
 function SpatialBatchNormalization:createIODescriptors(input)
    assert(input:dim() == 4)
    assert(torch.typename(self.weight) == 'torch.CudaTensor' and torch.typename(self.bias) == 'torch.CudaTensor',