From 93337d3e4470c344797a8838f2e6ec45e2506922 Mon Sep 17 00:00:00 2001
From: leozhang <nowerzt@gmail.com>
Date: Sat, 6 Apr 2024 21:57:42 +0800
Subject: [PATCH] Add more hyperparams for EvolutionOptimizer

---
 nncf/config/schemata/algo/filter_pruning.py   | 24 +++++++++++++++++++
 nncf/config/schemata/defaults.py              |  4 ++++
 .../filter_pruning/global_ranking/legr.py     | 19 ++++++++++++++-
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/nncf/config/schemata/algo/filter_pruning.py b/nncf/config/schemata/algo/filter_pruning.py
index 31c184346ac..d42e9da1e57 100644
--- a/nncf/config/schemata/algo/filter_pruning.py
+++ b/nncf/config/schemata/algo/filter_pruning.py
@@ -33,6 +33,10 @@
 from nncf.config.schemata.defaults import PRUNING_SCHEDULE
 from nncf.config.schemata.defaults import PRUNING_STEPS
 from nncf.config.schemata.defaults import PRUNING_TARGET
+from nncf.config.schemata.defaults import PRUNING_LEGR_POPULATION_SIZE
+from nncf.config.schemata.defaults import PRUNING_LEGR_NUM_SAMPLES
+from nncf.config.schemata.defaults import PRUNING_LEGR_MUTATE_PERCENT
+from nncf.config.schemata.defaults import PRUNING_LEGR_SIGMA_SCALE
 
 FILTER_PRUNING_SCHEDULE_OPTIONS = ["exponential", "exponential_with_bias", "baseline"]
 FILTER_IMPORTANCE_OPTIONS = ["L2", "L1", "geometric_median"]
@@ -162,6 +166,26 @@
                             description="Random seed for LeGR coefficients generation.",
                             default=PRUNING_LEGR_RANDOM_SEED,
                         ),
+                        "population_size": with_attributes(
+                            NUMBER,
+                            description="Size of population for the evolution algorithm.",
+                            default=PRUNING_LEGR_POPULATION_SIZE,
+                        ),
+                        "num_samples": with_attributes(
+                            NUMBER,
+                            description="Number of samples for the evolution algorithm.",
+                            default=PRUNING_LEGR_NUM_SAMPLES,
+                        ),
+                        "mutate_percent": with_attributes(
+                            NUMBER,
+                            description="Percent of mutate for the evolution algorithm.",
+                            default=PRUNING_LEGR_MUTATE_PERCENT,
+                        ),
+                        "scale_sigma": with_attributes(
+                            NUMBER,
+                            description="Scale sigma for the evolution algorithm.",
+                            default=PRUNING_LEGR_SIGMA_SCALE,
+                        ),
                     },
                     "additionalProperties": False,
                 },
diff --git a/nncf/config/schemata/defaults.py b/nncf/config/schemata/defaults.py
index d1a6471a768..3bf599fce95 100644
--- a/nncf/config/schemata/defaults.py
+++ b/nncf/config/schemata/defaults.py
@@ -59,6 +59,10 @@
 PRUNING_LEGR_TRAIN_STEPS = 200
 PRUNING_LEGR_MAX_PRUNING = 0.8
 PRUNING_LEGR_RANDOM_SEED = 42
+PRUNING_LEGR_POPULATION_SIZE = 64
+PRUNING_LEGR_NUM_SAMPLES = 16
+PRUNING_LEGR_MUTATE_PERCENT = 0.1
+PRUNING_LEGR_SIGMA_SCALE = 1
 
 SPARSITY_INIT = 0.0
 MAGNITUDE_SPARSITY_WEIGHT_IMPORTANCE = "normed_abs"
diff --git a/nncf/torch/pruning/filter_pruning/global_ranking/legr.py b/nncf/torch/pruning/filter_pruning/global_ranking/legr.py
index 2949ded0469..e1173505f0e 100644
--- a/nncf/torch/pruning/filter_pruning/global_ranking/legr.py
+++ b/nncf/torch/pruning/filter_pruning/global_ranking/legr.py
@@ -17,6 +17,10 @@
 from nncf.config.schemata.defaults import PRUNING_LEGR_MAX_PRUNING
 from nncf.config.schemata.defaults import PRUNING_LEGR_RANDOM_SEED
 from nncf.config.schemata.defaults import PRUNING_LEGR_TRAIN_STEPS
+from nncf.config.schemata.defaults import PRUNING_LEGR_POPULATION_SIZE
+from nncf.config.schemata.defaults import PRUNING_LEGR_NUM_SAMPLES
+from nncf.config.schemata.defaults import PRUNING_LEGR_MUTATE_PERCENT
+from nncf.config.schemata.defaults import PRUNING_LEGR_SIGMA_SCALE
 from nncf.torch.pruning.filter_pruning.global_ranking.evolutionary_optimization import EvolutionOptimizer
 from nncf.torch.pruning.filter_pruning.global_ranking.evolutionary_optimization import LeGREvolutionEnv
 from nncf.torch.pruning.filter_pruning.global_ranking.evolutionary_optimization import LeGRPruner
@@ -38,6 +42,10 @@ def __init__(
         generations: int = PRUNING_LEGR_GENERATIONS,
         max_pruning: float = PRUNING_LEGR_MAX_PRUNING,
         random_seed: int = PRUNING_LEGR_RANDOM_SEED,
+        population_size: int = PRUNING_LEGR_POPULATION_SIZE,
+        num_samples: int = PRUNING_LEGR_NUM_SAMPLES,
+        mutate_percent: float = PRUNING_LEGR_MUTATE_PERCENT,
+        scale_sigma: float = PRUNING_LEGR_SIGMA_SCALE
     ):
         """
         Initializing all necessary structures for optimization- LeGREvolutionEnv environment and EvolutionOptimizer
@@ -53,10 +61,19 @@ def __init__(
         self.num_generations = generations
         self.max_pruning = max_pruning
         self.train_steps = train_steps
+        self.population_size = population_size
+        self.num_samples = num_samples
+        self.mutate_percent = mutate_percent
+        self.scale_sigma = scale_sigma
 
         self.pruner = LeGRPruner(pruning_ctrl, target_model)
         init_filter_norms = self.pruner.init_filter_norms
-        agent_hparams = {"num_generations": self.num_generations}
+        agent_hparams = {"num_generations": self.num_generations, 
+                         "population_size": self.population_size,
+                         "num_samples": self.num_samples,
+                         "mutate_percent": self.mutate_percent,
+                         "sigma_scale": self.scale_sigma
+                         }
         self.agent = EvolutionOptimizer(init_filter_norms, agent_hparams, random_seed)
         self.env = LeGREvolutionEnv(
             self.pruner,