add hpo_params to all configs, make hpo return best_trials in case of multiobjective optimization

Author: nroope

src/pquant/configs/config_ap.yaml

@@ -1,5 +1,5 @@
 pruning_parameters:
-  disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true
+  disable_pruning_for_layers: []
   enable_pruning: true
   pruning_method: activation_pruning
   threshold: 0.2
@@ -46,3 +46,12 @@ training_parameters:
   rewind: never
   rounds: 1
   save_weights_epoch: -1
+hpo_parameters:
+    experiment_name: experiment_name
+    model_name: jet_tagger
+    num_trials: 1
+    sampler:
+      type: RandomSampler
+    hyperparameter_search:
+      numerical: {}
+      categorical: {}

src/pquant/configs/config_autosparse.yaml

@@ -49,3 +49,12 @@ training_parameters:
   rewind: never
   rounds: 1
   save_weights_epoch: -1.0
+hpo_parameters:
+    experiment_name: experiment_name
+    model_name: jet_tagger
+    num_trials: 1
+    sampler:
+      type: RandomSampler
+    hyperparameter_search:
+      numerical: {}
+      categorical: {}

src/pquant/configs/config_cs.yaml

@@ -45,3 +45,12 @@ training_parameters:
   rewind: post-training-stage
   rounds: 3
   save_weights_epoch: 2
+hpo_parameters:
+    experiment_name: experiment_name
+    model_name: jet_tagger
+    num_trials: 1
+    sampler:
+      type: RandomSampler
+    hyperparameter_search:
+      numerical: {}
+      categorical: {}

src/pquant/configs/config_fitcompress.yaml

@@ -44,3 +44,12 @@ training_parameters:
   rewind: never
   rounds: 1
   save_weights_epoch: -1
+hpo_parameters:
+    experiment_name: experiment_name
+    model_name: jet_tagger
+    num_trials: 1
+    sampler:
+      type: RandomSampler
+    hyperparameter_search:
+      numerical: {}
+      categorical: {}

src/pquant/configs/config_mdmm.yaml

@@ -55,6 +55,15 @@ fitcompress_parameters:
   greedy_astar : true
   approximate : true
   f_lambda : 1
+hpo_parameters:
+    experiment_name: experiment_name
+    model_name: jet_tagger
+    num_trials: 1
+    sampler:
+      type: RandomSampler
+    hyperparameter_search:
+      numerical: {}
+      categorical: {}
 
 # Note:
 # use_grad: true is having some bug... flip gradient impl not working as intended

src/pquant/configs/config_pdp.yaml

@@ -46,3 +46,12 @@ training_parameters:
   rewind: never
   rounds: 1
   save_weights_epoch: -1
+hpo_parameters:
+    experiment_name: experiment_name
+    model_name: jet_tagger
+    num_trials: 1
+    sampler:
+      type: RandomSampler
+    hyperparameter_search:
+      numerical: {}
+      categorical: {}

src/pquant/configs/config_wanda.yaml

@@ -49,3 +49,12 @@ training_parameters:
   rewind: never
   rounds: 1
   save_weights_epoch: -1
+hpo_parameters:
+    experiment_name: experiment_name
+    model_name: jet_tagger
+    num_trials: 1
+    sampler:
+      type: RandomSampler
+    hyperparameter_search:
+      numerical: {}
+      categorical: {}

src/pquant/core/hyperparameter_optimization.py

@@ -343,8 +343,10 @@ def run_optimization(self, model, **kwargs):
             n_trials=num_trials,
             n_jobs=1,
         )
-
-        return study.best_params
+        if len(self.objectives.keys()) == 1:
+            return study.best_params
+        else:
+            return study.best_trials
 
 
 def ap_config():

src/pquant/core/keras/layers.py

@@ -2251,7 +2251,7 @@ def post_training_prune(model, config, calibration_data):
     return apply_final_compression(model, config)
 
 
-def get_ebops(model):
+def get_ebops(model, **kwargs):
     ebops = 0
     for m in model.layers:
         if isinstance(m, (PQWeightBiasBase)):

src/pquant/core/torch/layers.py

@@ -1733,7 +1733,7 @@ def post_training_prune(model, config, calibration_data):
     return remove_compression_layers(model, config)
 
 
-def get_ebops(model):
+def get_ebops(model, **kwargs):
     ebops = 0
     for m in model.modules():
         if isinstance(m, (PQWeightBiasBase)):