re-implementation of Freeze BN and Rang for quantize

3 files changed, 24 insertions, 27 deletions

diff --git a/xmmdet/apis/train.py b/xmmdet/apis/train.py
index 785664db..a47cd082 100644
--- a/xmmdet/apis/train.py
+++ b/xmmdet/apis/train.py
@@ -10,7 +10,7 @@ from mmcv.utils import build_from_cfg
 from mmdet.core import DistEvalHook, EvalHook
 from mmdet.datasets import build_dataloader, build_dataset
 from ..utils import (get_root_logger, XMMDetEpochBasedRunner, XMMDetNoOptimizerHook, \
-                XMMDetDataParallel)
+                XMMDetDataParallel, FreezeRangeHook)
 
 
 def set_random_seed(seed, deterministic=False):
@@ -84,15 +84,13 @@ def train_detector(model,
 
     # build runner
     quantize = cfg.get('quantize', False)
-    freeze_range = bool(quantize)
     optimizer = build_optimizer(model, cfg.optimizer)
     runner = XMMDetEpochBasedRunner(
         model,
         optimizer=optimizer,
         work_dir=cfg.work_dir,
         logger=logger,
-        meta=meta,
-        freeze_range=freeze_range)
+        meta=meta)
     # an ugly workaround to make .log and .log.json filenames the same
     runner.timestamp = timestamp
 
@@ -115,6 +113,12 @@ def train_detector(model,
     if distributed:
         runner.register_hook(DistSamplerSeedHook())
 
+    # register train hooks
+    freeze_range = bool(quantize)
+    if freeze_range:
+        runner.register_hook(FreezeRangeHook())
+    #
+
     # register eval hooks
     if validate:
         # Support batch_size > 1 in validation
diff --git a/xmmdet/utils/__init__.py b/xmmdet/utils/__init__.py
index fe27d9f5..8b6923cd 100644
--- a/xmmdet/utils/__init__.py
+++ b/xmmdet/utils/__init__.py
@@ -1,7 +1,7 @@
 from mmdet.utils import *
 from .flops_counter import get_model_complexity_info
 from .logger import LoggerStream, get_root_logger
-from .runner import XMMDetEpochBasedRunner, XMMDetNoOptimizerHook, \
+from .runner import XMMDetEpochBasedRunner, XMMDetNoOptimizerHook, FreezeRangeHook, \
     mmdet_load_checkpoint, mmdet_save_checkpoint
 from .save_model import save_model_proto
 from .quantize import XMMDetQuantTrainModule, XMMDetQuantCalibrateModule, \
diff --git a/xmmdet/utils/runner.py b/xmmdet/utils/runner.py
index 52edbc56..07195703 100644
--- a/xmmdet/utils/runner.py
+++ b/xmmdet/utils/runner.py
@@ -1,7 +1,7 @@
 import torch
 import mmcv
 from mmcv.runner import EpochBasedRunner
-from mmcv.runner import OptimizerHook
+from mmcv.runner import OptimizerHook, HOOKS, Hook
 from pytorch_jacinto_ai import xnn
 from .quantize import is_mmdet_quant_module
 
@@ -21,27 +21,6 @@ def mmdet_save_checkpoint(model, *args, **kwargs):
 
 
 class XMMDetEpochBasedRunner(EpochBasedRunner):
-    def __init__(self, *args, **kwargs):
-        freeze_range = kwargs.pop('freeze_range', False)
-        super().__init__(*args, **kwargs)
-        self.freeze_range = freeze_range
-
-    def train(self, data_loader, **kwargs):
-        if self.freeze_range:
-            # currently we don't have a parameter that indicates whether we are doing QAT or not.
-            # Let us do it for all cases of training for the time being.
-            freeze_bn_epoch = (self.max_epochs//2)-1
-            freeze_range_epoch = (self.max_epochs//2)+1
-            if self.epoch > 0 and self.epoch >= freeze_bn_epoch:
-                xnn.utils.freeze_bn(self.model)
-            #
-            if self.epoch > 1 and self.epoch >= freeze_range_epoch:
-                xnn.layers.freeze_quant_range(self.model)
-            #
-        #
-        super().train(data_loader, **kwargs)
-
-
     def _get_model_orig(self):
         model_orig = self.model
         is_model_orig = True
@@ -84,9 +63,23 @@ class XMMDetEpochBasedRunner(EpochBasedRunner):
         #
 
 
+@HOOKS.register_module()
 class XMMDetNoOptimizerHook(OptimizerHook):
     def after_train_iter(self, runner):
         pass
 
 
+@HOOKS.register_module()
+class FreezeRangeHook(Hook):
+    def before_train_epoch(self, runner):
+        freeze_bn_epoch = (runner.max_epochs // 2) - 1
+        freeze_range_epoch = (runner.max_epochs // 2) + 1
+        if runner.epoch >= 1 and runner.epoch >= freeze_bn_epoch:
+            xnn.utils.print_once('Freezing BN')
+            xnn.utils.freeze_bn(runner.model)
+        #
+        if runner.epoch >= 2 and runner.epoch >= freeze_range_epoch:
+            xnn.utils.print_once('Freezing Activation ranges')
+            xnn.layers.freeze_quant_range(runner.model)
+        #