StructuredPruningModifier Refactor (#560)

* StructuredPruningModifier Refactor

* migrate packages

* style

* nit fix

* rebase test fixes

* rebase fixes + mask type for mfac

* Fix: test MFAC mask type

* Update src/sparseml/pytorch/sparsification/pruning/modifier_pruning_structured.py

* style

Author: bfineran

src/sparseml/pytorch/optim/__init__.py

@@ -36,7 +36,6 @@
 from .modifier_pruning import *
 from .modifier_quantization import *
 from .modifier_regularizer import *
-from .modifier_thinning import *
 from .optimizer import *
 from .sensitivity_as import *
 from .sensitivity_lr import *

src/sparseml/pytorch/optim/modifier_pruning.py

@@ -30,7 +30,6 @@
 from sparseml.pytorch.nn import Identity
 from sparseml.pytorch.optim.analyzer_pruning import ModulePruningAnalyzer
 from sparseml.pytorch.optim.mask_creator_pruning import (
-    DimensionSparsityMaskCreator,
     PruningMaskCreator,
     load_mask_creator,
 )
@@ -65,7 +64,6 @@
     "LegacyGMPruningModifier",
     "MovementPruningModifier",
     "LayerPruningModifier",
-    "StructuredPruningModifier",
 ]
 
 
@@ -874,191 +872,15 @@ def score_type(self) -> str:
         return self._score_type
 
 
-@PyTorchModifierYAML()
-class StructuredPruningModifier(_GMPruningModifier):
-    """
-    Gradually applies structured kernel sparsity to a given parameter or parameters
-    from init_sparsity until final_sparsity is reached over a given amount of time
-    and applied with an interpolated function for each step taken. Channel and filter
-    pruning supported.
-
-    A param_group_dependency_map must be provided that maps
-    groups of prunable parameter names that should have their dimensions pruned
-    together to a list of module parameter names that should be updated accordingly
-    when those parameters are pruned.
-
-    | Sample yaml:
-    |   !StructuredPruningModifier
-    |       param_groups: [
-    |           ["param.1.name","param.2.name"], ["param.3.name", "param.4.name"]
-    |       ]
-    |       init_sparsity: 0.05
-    |       final_sparsity: 0.8
-    |       start_epoch: 0.0
-    |       end_epoch: 10.0
-    |       update_frequency: 1.0
-    |       params: __ALL_PRUNABLE__
-    |       leave_enabled: True
-    |       inter_func: cubic
-    |       log_types: __ALL__
-    |       mask_type: filter
-    |       score_type: magnitude
-
-    :param init_sparsity: the initial sparsity for the param to start with at
-        start_epoch
-    :param final_sparsity: the final sparsity for the param to end with at end_epoch
-    :param start_epoch: The epoch to start the modifier at
-    :param end_epoch: The epoch to end the modifier at
-    :param update_frequency: The number of epochs or fraction of epochs to update at
-        between start and end
-    :param param_groups: list of list of parameter names that should be pruned together
-        during structured pruning so that their same indices may be removed. May be
-        useful for structures such as residual blocks or grouped convolutions. Can be
-        generated from an onnx export of the target module with
-        sparseml.onnx.optim.get_param_structured_pruning_group_dependencies by
-        splitting its comma separated keys into lists.
-        i.e. [["param.1.name","param.2.name"], ["param.3.name", "param.4.name"]]
-    :param params: A list of full parameter names or regex patterns of names to apply
-        pruning to.  Regex patterns must be specified with the prefix 're:'. __ALL__
-        will match to all parameters. __ALL_PRUNABLE__ will match to all ConvNd
-        and Linear layers' weights. Defualt is __ALL_PRUNABLE__
-    :param leave_enabled: True to continue masking the weights after end_epoch,
-        False to stop masking. Should be set to False if exporting the result
-        immediately after or doing some other prune
-    :param inter_func: the type of interpolation function to use:
-        [linear, cubic, inverse_cubic]
-    :param phased: True to enable a phased approach where pruning will
-        turn on and off with the update_frequency. Starts with pruning on
-        at start_epoch, off at start_epoch + update_frequency, and so on.
-    :param log_types: The loggers to allow the learning rate to be logged to,
-        default is __ALL__
-    :param mask_type: String to define type of structured sparsity (options: [
-        'channel', 'filter']), or a DimensionSparsityMaskCreator object.
-        default is 'filter'
-    :param score_type: Method used to score parameters for masking, i.e.
-        'magnitude', 'movement'. Default is 'magnitude'
-    """
-
-    def __init__(
-        self,
-        init_sparsity: float,
-        final_sparsity: float,
-        start_epoch: float,
-        end_epoch: float,
-        update_frequency: float,
-        param_groups: List[List[str]] = None,
-        params: Union[str, List[str]] = ALL_PRUNABLE_TOKEN,
-        leave_enabled: bool = True,
-        inter_func: str = "cubic",
-        phased: bool = False,
-        log_types: Union[str, List[str]] = ALL_TOKEN,
-        mask_type: Union[str, DimensionSparsityMaskCreator] = "filter",
-        score_type: str = "magnitude",
-    ):
-        if not isinstance(mask_type, DimensionSparsityMaskCreator) and (
-            mask_type not in ["channel", "filter"]
-        ):
-            raise ValueError(
-                "StructuredPruningModifier mask_type must be a "
-                "DimensionSparsityMaskCreator or designate 'channel' or 'filter' "
-                f"found {mask_type}"
-            )
-        super().__init__(
-            init_sparsity=init_sparsity,
-            final_sparsity=final_sparsity,
-            start_epoch=start_epoch,
-            end_epoch=end_epoch,
-            update_frequency=update_frequency,
-            params=params,
-            leave_enabled=leave_enabled,
-            inter_func=inter_func,
-            phased=phased,
-            log_types=log_types,
-            mask_type=mask_type,
-            global_sparsity=False,
-            score_type=score_type,
-        )
-
-        self._param_groups = param_groups or []
-
-    @BaseGMPruningModifier.sparsification_types.getter
-    def sparsification_types(self) -> List[SparsificationTypes]:
-        """
-        :return: the sparsification types this modifier instance will apply
-        """
-        return [SparsificationTypes.pruning, SparsificationTypes.structured]
-
-    @ModifierProp()
-    def param_groups(self) -> List[List[str]]:
-        """
-        :return: list of list of parameter names that should be pruned together
-            during structured pruning so that their same indices may be removed. May be
-            useful for structures such as residual blocks or grouped convolutions
-        """
-        return self._param_groups
-
-    @ModifierProp(serializable=False)
-    def global_sparsity(self) -> bool:
-        """
-        :return: True if global pruning is enabled, False otherwise
-        """
-        return self._global_sparsity
-
-    def _create_pruning_mask(
-        self, layers: List[Module], layer_names: List[str], param_names: List[str]
-    ) -> ModuleParamPruningMask:
-        # find and validate parameter groups for structured pruning
-        full_param_names = [
-            f"{layer_name}.{param_name}"
-            for layer_name, param_name in zip(layer_names, param_names)
-        ]
-        param_name_to_idx = dict(zip(full_param_names, range(len(full_param_names))))
-        param_group_idxs = []
-        added_idxs = set()
-
-        for param_group in self._param_groups:
-            group_idxs = []
-            for param_name in param_group:
-                if param_name not in param_name_to_idx:
-                    raise ValueError(
-                        f"param {param_name} from param_groups "
-                        f"not found in pruning modifier params {full_param_names}"
-                    )
-                param_idx = param_name_to_idx[param_name]
-                if param_idx in added_idxs:
-                    raise ValueError(
-                        "found repeated param name in param_groups " f"{param_name}"
-                    )
-                group_idxs.append(param_idx)
-                added_idxs.add(param_idx)
-            param_group_idxs.append(group_idxs)
-        for idx in range(len(full_param_names)):
-            if idx not in added_idxs:
-                param_group_idxs.append([idx])
-
-        self._mask_creator.set_tensor_group_idxs(param_group_idxs)
-
-        return ModuleParamPruningMask(
-            layers,
-            param_names,
-            layer_names=layer_names,
-            mask_creator=self._mask_creator,
-            global_sparsity=False,
-            score_type=self._score_type,
-        )
-
-
 @PyTorchModifierYAML()
 class LayerPruningModifier(ScheduledUpdateModifier):
     """
     Class for pruning away layers within a module
     (replaces with sparseml.pytorch.nn.Identity).
-
     | Sample yaml:
     |   !LayerPruningModifier
     |       layers: ['bert.encoder.layer.6', 'bert.encoder.layer.7']
     |
-
     :param layers: A list of full layer names to apply pruning to.
         __ALL_ will match to all layers. __ALL_PRUNABLE__ will match to all ConvNd
         and Linear layers
@@ -1124,7 +946,6 @@ def initialize(
     ):
         """
         Grab the layers and apply if epoch in range to control pruning for.
-
         :param module: the PyTorch model/module to modify
         :param epoch: The epoch to initialize the modifier and module at.
             Defaults to 0 (start of the training process)
@@ -1145,7 +966,6 @@ def finalize(
     ):
         """
         Cleans up any remaining hooks
-
         :param module: The model/module to finalize the modifier for.
             Marked optional so state can still be cleaned up on delete,
             but generally should always be passed in.
@@ -1165,7 +985,6 @@ def update(
     ):
         """
         Update to enable and disable the layers when chosen.
-
         :param module: module to modify
         :param optimizer: optimizer to modify
         :param epoch: current epoch and progress within the current epoch
@@ -1180,7 +999,6 @@ def log_update(
     ):
         """
         Check whether to log an update for the state of the modifier.
-
         :param module: module to modify
         :param optimizer: optimizer to modify
         :param epoch: current epoch and progress within the current epoch

src/sparseml/pytorch/sparsification/__init__.py

@@ -22,4 +22,5 @@
 # flake8: noqa
 
 from .info import *
+from .modifier_thinning import *
 from .pruning import *

…rseml/pytorch/optim/modifier_thinning.py …orch/sparsification/modifier_thinning.pysrc/sparseml/pytorch/optim/modifier_thinning.py renamed to src/sparseml/pytorch/sparsification/modifier_thinning.py src/sparseml/pytorch/optim/modifier_thinning.py renamed to src/sparseml/pytorch/sparsification/modifier_thinning.py

@@ -26,4 +26,5 @@
 from .modifier_pruning_constant import *
 from .modifier_pruning_magnitude import *
 from .modifier_pruning_mfac import *
+from .modifier_pruning_structured import *
 from .scorer import *

src/sparseml/pytorch/sparsification/pruning/__init__.py

@@ -217,8 +217,12 @@ def _get_scorer(self, params: List[Parameter]) -> PruningParamsScorer:
         """
         return MagnitudePruningParamsScorer(params)
 
-    def _get_mask_creator(self) -> PruningMaskCreator:
+    def _get_mask_creator(
+        self, param_names: List[str], params: List[Parameter]
+    ) -> PruningMaskCreator:
         """
+        :param names: full names of parameters to be pruned
+        :param params: list of Parameters to be masked
         :return: mask creator object to be used by this pruning algorithm
         """
         if self._mask_type == "unstructured":

src/sparseml/pytorch/sparsification/pruning/modifier_pruning_acdc.py

@@ -174,8 +174,12 @@ def sparsification_types(self) -> List[SparsificationTypes]:
         return [SparsificationTypes.pruning]
 
     @abstractmethod
-    def _get_mask_creator(self) -> PruningMaskCreator:
+    def _get_mask_creator(
+        self, param_names: List[str], params: List[Parameter]
+    ) -> PruningMaskCreator:
         """
+        :param names: full names of parameters to be pruned
+        :param params: list of Parameters to be masked
         :return: mask creator object to be used by this pruning algorithm
         """
         raise NotImplementedError()
@@ -241,7 +245,7 @@ def mask_creator(self) -> Optional[PruningMaskCreator]:
         """
         :return: mask creator object used by this pruning algorithm
         """
-        raise self._mask_creator
+        return self._mask_creator
 
     @property
     def scorer(self) -> Optional[PruningParamsScorer]:
@@ -295,13 +299,15 @@ def initialize(
         layer_names = [nlp.layer_name for nlp in named_layers_and_params]
 
         # initialize mask_creator and scorer
-        self._mask_creator = self._get_mask_creator()
-        self._scorer = self._get_scorer(
-            params=[
-                getattr(layer, param_name)
-                for layer, param_name in zip(layers, param_names)
-            ]
-        )
+        params = [
+            getattr(layer, param_name) for layer, param_name in zip(layers, param_names)
+        ]
+        full_param_names = [
+            f"{layer_name}.{param_name}"
+            for layer_name, param_name in zip(layer_names, param_names)
+        ]
+        self._mask_creator = self._get_mask_creator(full_param_names, params)
+        self._scorer = self._get_scorer(params)
 
         self._module_masks = self._create_pruning_mask(layers, layer_names, param_names)
         self._analyzers = self._create_analyzers(layers, layer_names, param_names)

src/sparseml/pytorch/sparsification/pruning/modifier_pruning_base.py

@@ -21,7 +21,7 @@
 
 import torch
 from torch import Tensor
-from torch.nn import Module
+from torch.nn import Module, Parameter
 
 from sparseml.pytorch.optim.modifier import (
     ModifierProp,
@@ -135,8 +135,12 @@ def __init__(
             parent_class_kwarg_names=["params"],
         )
 
-    def _get_mask_creator(self) -> PruningMaskCreator:
+    def _get_mask_creator(
+        self, param_names: List[str], params: List[Parameter]
+    ) -> PruningMaskCreator:
         """
+        :param names: full names of parameters to be pruned
+        :param params: list of Parameters to be masked
         :return: mask creator object to be used by this pruning algorithm
         """
         return ConstantMaskCreator()

src/sparseml/pytorch/sparsification/pruning/modifier_pruning_constant.py

@@ -135,11 +135,21 @@ def __init__(
             end_comparator=-1,
             global_sparsity=self._use_global_sparsity,
             allow_reintroduction=False,
-            parent_class_kwarg_names=["init_sparsity", "final_sparsity", "params"],
+            parent_class_kwarg_names=[
+                "init_sparsity",
+                "final_sparsity",
+                "params",
+                "leave_enabled",
+                "mask_type",
+            ],
         )
 
-    def _get_mask_creator(self) -> PruningMaskCreator:
+    def _get_mask_creator(
+        self, param_names: List[str], params: List[Parameter]
+    ) -> PruningMaskCreator:
         """
+        :param names: full names of parameters to be pruned
+        :param params: list of parameters to be masked
         :return: mask creator object to be used by this pruning algorithm
         """
         if self.mask_type == "unstructured":