Added optimization level for new API & support sub-parameter setting for strategy (#222)

Signed-off-by: yiliu30 <[email protected]>

Author: yiliu30

neural_compressor/conf/config.py

@@ -841,6 +841,7 @@ def percent_to_float(data):
         },
     },
     Optional('use_bf16', default=True): bool,
+    Optional('optimization_level', default=1): And(int, lambda level: level in [0, 1]),
     Optional('graph_optimization'): graph_optimization_schema,
     Optional('mixed_precision'): mixed_precision_schema,
 
@@ -1111,6 +1112,7 @@ def percent_to_float(data):
                                                      'activation': {}},
                                     }): dict,
     Optional('use_bf16', default=False): bool,
+    Optional('optimization_level', default=1): int,
     Optional('tuning', default={
         'strategy': {'name': 'basic'},
         'accuracy_criterion': {'relative': 0.01, 'higher_is_better': True},
@@ -1346,8 +1348,17 @@ def map_pyconfig_to_cfg(self, pythonic_config):
                 'tuning.exit_policy.max_trials': pythonic_config.quantization.max_trials,
                 'tuning.exit_policy.performance_only': pythonic_config.quantization.performance_only,
                 'use_bf16': pythonic_config.quantization.use_bf16,
+                'quantization.optimization_level': pythonic_config.quantization.optimization_level,
                 'reduce_range': pythonic_config.quantization.reduce_range
             })
+            if pythonic_config.quantization.strategy_kwargs:
+                st_kwargs = pythonic_config.quantization.strategy_kwargs
+                for st_key in ['sigopt_api_token', 'sigopt_project_id', 'sigopt_experiment_name', \
+                    'accuracy_weight', 'latency_weight']:
+                    if st_key in st_kwargs:
+                        st_val =  st_kwargs[st_key]
+                        mapping.update({'tuning.strategy.' + st_key: st_val})
+            
         if pythonic_config.distillation is not None:
             mapping.update({
                 'distillation.train.criterion': pythonic_config.distillation.criterion,

neural_compressor/conf/pythonic_config.py

@@ -34,17 +34,33 @@ def __init__(self,
                  op_type_list=None,
                  op_name_list=None,
                  strategy='basic',
+                 strategy_kwargs=None,
                  objective='performance',
                  timeout=0,
                  max_trials=100,
                  performance_only=False,
                  reduce_range=None,
                  use_bf16=True,
+                 optimization_level=1,
                  accuracy_criterion=accuracy_criterion):
         extra_precisions = ["bf16"] if use_bf16 else []
-        super().__init__(inputs, outputs, backend, device, calibration_sampling_size, op_type_list,
-                         op_name_list, strategy, objective, timeout, max_trials, performance_only,
-                         reduce_range, extra_precisions, accuracy_criterion)
+        super().__init__(inputs=inputs,
+                         outputs=outputs,
+                         backend=backend,
+                         device=device,
+                         calibration_sampling_size=calibration_sampling_size,
+                         op_type_list=op_type_list,
+                         op_name_list=op_name_list,
+                         strategy=strategy,
+                         strategy_kwargs=strategy_kwargs,
+                         objective=objective,
+                         timeout=timeout,
+                         max_trials=max_trials,
+                         performance_only=performance_only,
+                         reduce_range=reduce_range,
+                         extra_precisions=extra_precisions,
+                         optimization_level=optimization_level,
+                         accuracy_criterion=accuracy_criterion)
         self._approach = approach
 
     @property

neural_compressor/config.py

@@ -307,12 +307,14 @@ def __init__(self,
                  op_type_list=None,
                  op_name_list=None,
                  strategy="basic",
+                 strategy_kwargs=None,
                  objective="performance",
                  timeout=0,
                  max_trials=100,
                  performance_only=False,
                  reduce_range=None,
                  extra_precisions=["bf16"],
+                 optimization_level=1,
                  accuracy_criterion=accuracy_criterion):
         self._inputs = inputs
         self._outputs = outputs
@@ -321,13 +323,15 @@ def __init__(self,
         self._op_type_list = op_type_list
         self._op_name_list = op_name_list
         self._strategy = strategy
+        self._strategy_kwargs = strategy_kwargs
         self._objective = objective
         self._timeout = timeout
         self._max_trials = max_trials
         self._performance_only = performance_only
         self._reduce_range = reduce_range
         self._extra_precisions = extra_precisions \
             if isinstance(extra_precisions, List) else [extra_precisions]
+        self._optimization_level = optimization_level
         self.use_bf16 = "bf16" in self._extra_precisions
         self._accuracy_criterion = accuracy_criterion
         self._calibration_sampling_size = calibration_sampling_size
@@ -346,6 +350,14 @@ def extra_precisions(self, extra_precisions):
             self._extra_precisions = extra_precisions
             self._use_bf16 = "bf16" in extra_precisions
 
+    @property
+    def optimization_level(self):
+        return self._optimization_level
+    
+    @optimization_level.setter
+    def optimization_level(self, optimization_level):
+        self._optimization_level = optimization_level
+
     @property
     def reduce_range(self):
         return self._reduce_range
@@ -399,9 +411,17 @@ def strategy(self):
     @strategy.setter
     def strategy(self, strategy):
         if check_value('strategy', strategy, str,
-            ['basic', 'mse', 'bayesian', 'random', 'exhaustive']):
+            ['basic', 'mse', 'bayesian', 'random', 'exhaustive', 'sigopt', 'tpe']):
             self._strategy = strategy
 
+    @property
+    def strategy_kwargs(self):
+        return self._strategy_kwargs
+    
+    @strategy_kwargs.setter
+    def strategy_kwargs(self, strategy_kwargs):
+        self._strategy_kwargs = strategy_kwargs
+
     @property
     def op_name_list(self):
         return self._op_name_list
@@ -480,11 +500,12 @@ def inputs(self, inputs):
 
 
 class TuningCriterion:
-    def __init__(self, strategy="basic", timeout=0, max_trials=100, objective="performance"):
+    def __init__(self, strategy="basic", strategy_kwargs=None, timeout=0, max_trials=100, objective="performance"):
         self._strategy = strategy
         self._timeout = timeout
         self._max_trials = max_trials
         self._objective = objective
+        self._strategy_kwargs = strategy_kwargs
 
     @property
     def max_trials(self):
@@ -521,9 +542,16 @@ def strategy(self):
     @strategy.setter
     def strategy(self, strategy):
         if check_value('strategy', strategy, str,
-            ['basic', 'mse', 'bayesian', 'random', 'exhaustive']):
+            ['basic', 'mse', 'bayesian', 'random', 'exhaustive', 'sigopt', 'tpe']):
             self._strategy = strategy
-
+    
+    @property
+    def strategy_kwargs(self):
+        return self._strategy_kwargs
+    
+    @strategy_kwargs.setter
+    def strategy_kwargs(self, strategy_kwargs):
+        self._strategy_kwargs = strategy_kwargs
 
 tuning_criterion = TuningCriterion()
 
@@ -540,6 +568,7 @@ def __init__(self,
                  op_name_list=None,
                  reduce_range=None,
                  extra_precisions = ["bf16"],
+                 optimization_level=1,
                  tuning_criterion=tuning_criterion,
                  accuracy_criterion=accuracy_criterion,
     ):
@@ -551,11 +580,13 @@ def __init__(self,
                          op_type_list=op_type_list,
                          op_name_list=op_name_list,
                          strategy=tuning_criterion.strategy,
+                         strategy_kwargs=tuning_criterion.strategy_kwargs,
                          objective=tuning_criterion.objective,
                          timeout=tuning_criterion.timeout,
                          max_trials=tuning_criterion.max_trials,
                          reduce_range=reduce_range,
                          extra_precisions=extra_precisions,
+                         optimization_level=optimization_level,
                          accuracy_criterion=accuracy_criterion)
         self.approach = approach
 
@@ -578,10 +609,12 @@ def __init__(self,
                  op_type_list=None,
                  op_name_list=None,
                  reduce_range=None,
-                 extra_precisions=["bf16"]):
+                 extra_precisions=["bf16"],
+                 optimization_level=1):
         super().__init__(inputs=inputs, outputs=outputs, device=device, backend=backend,
                          op_type_list=op_type_list, op_name_list=op_name_list,
-                         reduce_range=reduce_range, extra_precisions=extra_precisions)
+                         reduce_range=reduce_range, extra_precisions=extra_precisions, 
+                         optimization_level=optimization_level)
         self._approach = 'quant_aware_training'
 
     @property

test/strategy/test_basic.py

@@ -155,7 +155,7 @@ def build_fake_model():
             tf.import_graph_def(graph_def, name='')
     return graph
 
-class TestQuantization(unittest.TestCase):
+class TestBasicTuningStrategy(unittest.TestCase):
 
     @classmethod
     def setUpClass(self):
@@ -217,6 +217,25 @@ def test_run_basic_max_trials_multimetric_weight(self):
         quantizer.model = self.constant_graph
         quantizer.fit()
 
+        
+    def test_run_basic_one_trial_new_api(self):
+        from neural_compressor.quantization import fit
+        from neural_compressor.config import AccuracyCriterion, AccuracyLoss, PostTrainingQuantConfig, TuningCriterion
+        from neural_compressor.data import DATASETS, DATALOADERS
+        
+        # dataset and dataloader
+        dataset = DATASETS("tensorflow")["dummy"](((100, 3, 3, 1)))
+        dataloader = DATALOADERS["tensorflow"](dataset)
+        
+        # tuning and accuracy criterion
+        tolerable_loss = AccuracyLoss(0.01)
+        accuracy_criterion = AccuracyCriterion(criterion='relative', tolerable_loss=tolerable_loss)
+        tuning_criterion = TuningCriterion(strategy='basic')
+        conf = PostTrainingQuantConfig(approach="static", backend="tensorflow", 
+                                       tuning_criterion=tuning_criterion,
+                                       accuracy_criterion=accuracy_criterion)
+        q_model = fit(model=self.constant_graph, conf=conf, calib_dataloader= dataloader, eval_dataloader=dataloader)
+        self.assertIsNotNone(q_model)
 
 if __name__ == "__main__":
     unittest.main()

test/strategy/test_sigopt.py

@@ -104,7 +104,7 @@ def build_fake_model():
     return graph
 
 @unittest.skipIf(CONDITION , "missing the env variables 'SIGOPT_API_TOKEN' or 'SIGOPT_PROJECT_ID'")
-class TestQuantization(unittest.TestCase):
+class TestSigoptTuningStrategy(unittest.TestCase):
 
     @classmethod
     def setUpClass(self):
@@ -140,6 +140,29 @@ def test_run_basic_max_trials(self):
         quantizer.eval_dataloader = common.DataLoader(dataset)
         quantizer.model = self.constant_graph
         quantizer.fit()
+        
+    def test_run_sigopt_one_trial_new_api(self):
+        from neural_compressor.quantization import fit
+        from neural_compressor.config import AccuracyCriterion, AccuracyLoss, PostTrainingQuantConfig, TuningCriterion
+        from neural_compressor.data import DATASETS, DATALOADERS
+        
+        # dataset and dataloader
+        dataset = DATASETS("tensorflow")["dummy"](((100, 3, 3, 1)))
+        dataloader = DATALOADERS["tensorflow"](dataset)
+        
+        # tuning and accuracy criterion
+        tolerable_loss = AccuracyLoss(0.01)
+        accuracy_criterion = AccuracyCriterion(criterion='relative', tolerable_loss=tolerable_loss)
+        strategy_kwargs = {'sigopt_api_token': 'sigopt_api_token_test', 
+                           'sigopt_project_id': 'sigopt_project_id_test',
+                           'sigopt_experiment_name': 'nc-tune'}
+        tuning_criterion = TuningCriterion(strategy='sigopt', strategy_kwargs=strategy_kwargs, max_trials=3)
+        conf = PostTrainingQuantConfig(approach="static", backend="tensorflow", 
+                                       tuning_criterion=tuning_criterion,
+                                       accuracy_criterion=accuracy_criterion)
+        q_model = fit(model=self.constant_graph, conf=conf, calib_dataloader= dataloader, eval_dataloader=dataloader)
+        self.assertIsNotNone(q_model)
+
 
 if __name__ == "__main__":
     unittest.main()