Browse Source

fix(mge/quant): avoid updating the scale and zero point in eval mode

GitOrigin-RevId: dfb08cf701
tags/v0.5.0
Megvii Engine Team 5 years ago
parent
commit
206521cd33
2 changed files with 9 additions and 1 deletions
  1. +2
    -1
      python_module/megengine/module/module.py
  2. +7
    -0
      python_module/megengine/quantization/observer.py

+ 2
- 1
python_module/megengine/module/module.py View File

@@ -294,7 +294,8 @@ class Module(metaclass=ABCMeta):
self.training = mode self.training = mode


def fn(x) -> None: def fn(x) -> None:
x.training = mode
if x is not self:
x.train(mode=mode)


self.apply(fn) self.apply(fn)




+ 7
- 0
python_module/megengine/quantization/observer.py View File

@@ -56,6 +56,13 @@ class Observer(Module):
def disable(self): def disable(self):
self.enabled = False self.enabled = False


def train(self, mode: bool = True) -> None:
super().train(mode)
if mode:
self.enable()
else:
self.disable()

@abstractmethod @abstractmethod
def forward(self, x): def forward(self, x):
pass pass


Loading…
Cancel
Save