Automatic set gmp_config (PaddlePaddle#1103)

* Automatic set gmp_config * Automatic set gmp_config * add a note. * add a note.
lizexu123 · May 18, 2022 · e6ae764 · e6ae764
1 parent b05085e
commit e6ae764
Showing 1 changed file with 17 additions and 1 deletion.
diff --git a/paddleslim/auto_compression/compressor.py b/paddleslim/auto_compression/compressor.py
@@ -248,6 +248,21 @@ def _prepare_program(self, program, feed_target_names, fetch_targets,
                                          feed_target_names, fetch_targets)
 
         config_dict = dict(config._asdict())
+        if config_dict["prune_strategy"] == "gmp" and config_dict[
+                'gmp_config'] is None:
+            _logger.info(
+                "Calculating the iterations per epoch……(It will take some time)")
+            # NOTE:XXX: This way of calculating the iters needs to be improved.
+            iters_per_epoch = len(list(self.train_dataloader()))
+            total_iters = self.train_config.epochs * iters_per_epoch
+            config_dict['gmp_config'] = {
+                'stable_iterations': 0,
+                'pruning_iterations': 0.45 * total_iters,
+                'tunning_iterations': 0.45 * total_iters,
+                'resume_iteration': -1,
+                'pruning_steps': 100,
+                'initial_ratio': 0.15,
+            }
         ### add prune program
         self._pruner = None
         if 'prune' in strategy:
@@ -280,13 +295,14 @@ def _prepare_program(self, program, feed_target_names, fetch_targets,
                 test_program_info)
         if self.train_config.sparse_model:
             from ..prune.unstructured_pruner import UnstructuredPruner
+            # NOTE: The initialization parameter of this pruner doesn't work, it is only used to call the 'set_static_masks' function
             self._pruner = UnstructuredPruner(
                 train_program_info.program,
                 mode='ratio',
                 ratio=0.75,
                 prune_params_type='conv1x1_only',
                 place=self._places)
-            self._pruner.set_static_masks()
+            self._pruner.set_static_masks()  # Fixed model sparsity
 
         self._exe.run(train_program_info.startup_program)