From 186bacfb7164b3776725a6af620aa4cf8d20628f Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Thu, 22 Jul 2021 14:10:03 +0800 Subject: [PATCH] fix(mge): recover bn freeze fastpath execution GitOrigin-RevId: 13e58c9fba9638ddb1820f4629488086185a8337 --- imperative/python/megengine/module/batchnorm.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/imperative/python/megengine/module/batchnorm.py b/imperative/python/megengine/module/batchnorm.py index 2e95f25a6..efeaba17d 100644 --- a/imperative/python/megengine/module/batchnorm.py +++ b/imperative/python/megengine/module/batchnorm.py @@ -100,6 +100,15 @@ class _BatchNorm(Module): if _bias is not None: _bias = _bias.detach() + # fastpath excution for freeze + scale = (self.running_var + self.eps) ** (-0.5) + if _weight is not None: + scale *= _weight + bias = -self.running_mean * scale + if _bias is not None: + bias += _bias + return inp * scale + bias + if self.training and self.track_running_stats: exponential_average_factor = self.momentum else: -- GitLab