[Fix] AdaRound load state dict missing alpha. (#140)

* [Fix] AdaRound load state dict missing alpha.

* [Feat] per channel hooks for buffer

Co-authored-by: fanyunqian <[email protected]>

Author: PannenetsF

mqbench/fake_quantize/adaround_quantizer.py

@@ -48,7 +48,7 @@ def __init__(self, observer, **observer_kwargs):
         self.register_buffer('scale', torch.tensor([1.0], dtype=torch.float))
         self.register_buffer('zero_point', torch.tensor([0], dtype=torch.int))
         self.adaround = False
-        self.load_state_dict_hook = PerChannelLoadHook(self)
+        self.load_state_dict_hook = PerChannelLoadHook(self, hook_param=['scale', 'zero_point', 'alpha'])
 
     def init(self, weight_tensor: torch.Tensor, round_mode='learned_hard_sigmoid', ):
         self.adaround = True

mqbench/utils/hook.py

@@ -31,16 +31,25 @@ def __call__(self, module, input_batch, output_batch):
 
 
 class PerChannelLoadHook:
-    def __init__(self, module):
+    def __init__(self, module, hook_param=["scale", "zero_point"]):
         self.hook = module._register_load_state_dict_pre_hook(partial(self.hook_fn, module=module))
+        self.hook_param = hook_param
 
     def hook_fn(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs,
                 module):
         if module.ch_axis == -1:
             # no per-channel parameters
             return
         for module_key, param in module._parameters.items():
-            if module_key not in ["scale", "zero_point"]:
+            if module_key not in self.hook_param:
+                continue
+            candidate = prefix + module_key
+            if candidate in state_dict:
+                input_param = state_dict[candidate]
+                if param.shape != input_param.shape:
+                    param.data = torch.ones_like(input_param, dtype=param.dtype, device=param.device)
+        for module_key, param in module._buffers.items():
+            if module_key not in self.hook_param:
                 continue
             candidate = prefix + module_key
             if candidate in state_dict: