提交 186bacfb 编写于 作者: M Megvii Engine Team

fix(mge): recover bn freeze fastpath execution

GitOrigin-RevId: 13e58c9fba9638ddb1820f4629488086185a8337
上级 5f558042
......@@ -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:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册