Add hawq_v2 tuning strategy (#230)

Signed-off-by: yiliu30 <[email protected]>
Co-authored-by: lvliang-intel <[email protected]>
Co-authored-by: chen, suyue <[email protected]>
Co-authored-by: xinhe <[email protected]>
Co-authored-by: Ray <[email protected]>
Signed-off-by: zehao-intel <[email protected]>

Author: 5 people

examples/.config/model_params_pytorch.json

@@ -9,6 +9,24 @@
       "batch_size": 100,
       "new_benchmark": false
     },
+      "efficientnet_b0_fx": {
+      "model_src_dir": "image_recognition/torchvision_models/quantization/ptq/cpu/fx/",
+      "dataset_location": "/tf_dataset/pytorch/ImageNet/raw",
+      "input_model": "",
+      "yaml": "conf.yaml",
+      "strategy": "hawq_v2",
+      "batch_size": 100,
+      "new_benchmark": false
+    },
+      "efficientnet_b3_fx": {
+      "model_src_dir": "image_recognition/torchvision_models/quantization/ptq/cpu/fx/",
+      "dataset_location": "/tf_dataset/pytorch/ImageNet/raw",
+      "input_model": "",
+      "yaml": "conf.yaml",
+      "strategy": "hawq_v2",
+      "batch_size": 100,
+      "new_benchmark": false
+      },
     "resnet18_fx": {
       "model_src_dir": "image_recognition/torchvision_models/quantization/ptq/cpu/fx/",
       "dataset_location": "/tf_dataset/pytorch/ImageNet/raw",

examples/pytorch/image_recognition/torchvision_models/quantization/ptq/cpu/fx/conf.yaml

@@ -77,4 +77,4 @@ tuning:
     relative:  0.01                                  # optional. default value is relative, other value is absolute. this example allows relative accuracy loss: 1%.
   exit_policy:
     timeout: 0                                       # optional. tuning timeout (seconds). default value is 0 which means early stop. combine with max_trials field to decide when to exit.
-  random_seed: 9527                                  # optional. random seed for deterministic tuning.
+  random_seed: 9527                                  # optional. random seed for deterministic tuning.

neural_compressor/adaptor/pytorch.py

@@ -30,7 +30,6 @@
 from .query import QueryBackendCapability
 from ..experimental.data.dataloaders.base_dataloader import BaseDataLoader
 
-
 torch = LazyImport("torch")
 json = LazyImport("json")
 hvd = LazyImport("horovod.torch")
@@ -1094,6 +1093,34 @@ def is_fused_module(self, module):
             return True
         else:
             return False
+        
+    def calculate_hessian_trace(self,
+                                fp32_model, 
+                                dataloader, 
+                                q_model,
+                                criterion, 
+                                enable_act = False
+                                ):
+        """Calculate hessian trace.
+
+        Args:
+            fp32_model: The original fp32 model.
+            criterion: The loss function for calculate the hessian trace. # loss = criterion(output, target)
+            dataloader: The dataloader for calculate the gradient.
+            q_model: The INT8 AMAP model.
+            enable_act: Enabling quantization error or not.
+            
+        Return:
+            hessian_trace(Dict[Tuple, float]), key: (op_name, op_type); value: hessian trace.
+        """
+        from .torch_utils.hawq_metric import hawq_top
+        op_to_traces=hawq_top(fp32_model=fp32_model,
+                              dataloader=dataloader,
+                              q_model=q_model,
+                              criterion=criterion,
+                              enable_act=enable_act)
+        return op_to_traces
+        pass
 
 
 unify_op_type_mapping = {