mask creator refactor (#599) (#601)

Author: bfineran

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

@@ -31,6 +31,7 @@
     "GroupedPruningMaskCreator",
     "UnstructuredPruningMaskCreator",
     "FourBlockMaskCreator",
+    "BlockMaskCreator",
 ]
 
 
@@ -422,3 +423,148 @@ def _map_mask_to_tensor(
             permute_val.insert(1, len(permute_val))
             block_mask = block_mask.permute(*permute_val)
         return block_mask
+
+
+class BlockMaskCreator(GroupedPruningMaskCreator):
+    """
+    Structured sparsity mask creator that groups the input tensor into blocks of
+    shape block_shape.
+
+    :param block_shape: The shape in and out channel should take in blocks.  Should be
+        a list of exactly two integers that divide the input tensors evenly on the
+        channel dimensions.  -1 for a dimension blocks across the entire dimension
+    :param grouping_fn_name: The name of the torch grouping function to reduce
+        dimensions by
+    """
+
+    def __init__(
+        self,
+        block_shape: List[int],
+        grouping_fn_name: str = "mean",
+    ):
+        if len(block_shape) < 2:
+            raise ValueError(
+                (
+                    "Invalid block_shape: {}, "
+                    "block_shape must have length == 2 for in and out channels"
+                ).format(block_shape)
+            )
+
+        if len(block_shape) > 2 and not all([shape == 1 for shape in block_shape[2:]]):
+            # after in and out channels, only 1 can be used for other dimensions
+            raise ValueError(
+                (
+                    "Invalid block_shape: {}, "
+                    "block_shape for indices not in [0, 1] must be equal to 1"
+                ).format(block_shape)
+            )
+
+        self._block_shape = block_shape
+        self._grouping_fn_name = grouping_fn_name
+
+    def group_tensor(self, tensor: Tensor) -> Tensor:
+        """
+        :param tensor: The tensor to transform
+        :return: The mean values of the tensor grouped by blocks of shape
+            self._block_shape
+        """
+        blocked_tens_shape = self._get_blocked_tens_shape_and_validate(tensor.shape)
+        blocked_tensor = tensor.reshape(blocked_tens_shape)
+        reduced_blocks = GroupedPruningMaskCreator.reduce_tensor(
+            blocked_tensor, 1, self._grouping_fn_name
+        )
+        return reduced_blocks.type(tensor.type())
+
+    def _map_mask_to_tensor(
+        self,
+        grouped_mask: Tensor,
+        original_tensor_shape: torch.Size,
+        tensor_idx: Optional[int] = None,
+    ) -> Tensor:
+        """
+        :param grouped_mask: A binary mask the size of a tensor from group_tensor
+        :param original_tensor_shape: Shape of the original tensor grouped_mask
+            derives from
+        :param tensor_idx: optional index this tensor was passed into a tensor
+            list for mask creation
+        :return: The values from grouped_mask mapped to a tensor of size
+            original_tensor_shape
+        """
+        blocked_tens_shape = self._get_blocked_tens_shape_and_validate(
+            original_tensor_shape
+        )
+        # expand so every element has a corresponding value in the original tensor
+        block_mask = grouped_mask.reshape(blocked_tens_shape[0], blocked_tens_shape[2])
+        block_mask = block_mask.unsqueeze(1)
+        block_mask = block_mask.expand(*blocked_tens_shape).contiguous()
+        return block_mask.reshape(original_tensor_shape)
+
+    def _get_blocked_tens_shape_and_validate(
+        self,
+        tens_shape: torch.Size,
+    ) -> List[int]:
+        """
+        :param tens_shape: The shape of the tensor to group in blocks
+        :return: shape of tens when blocked by block_shape
+        :raise: ValueError if we are unable to block tens by shape block_shape
+        """
+        block_shape = self._block_shape
+        n_dims = len(tens_shape)
+        while len(block_shape) < n_dims:  # Conv will have block shape [X, Y, 1, ..., 1]
+            block_shape.append(1)
+        for idx, shape in enumerate(block_shape):
+            if shape == -1:
+                block_shape[idx] = tens_shape[idx]
+        # Validate
+        if n_dims < 2:
+            raise ValueError(
+                "Invalid tensor shape {}."
+                " BlockSparsityMaskCreator can only create masks from tensors with 2 or"
+                " more dimensions, tensor has {}.".format(tens_shape, n_dims)
+            )
+        for tens_dim, block_dim in zip(tens_shape, block_shape):
+            if tens_dim % block_dim != 0:
+                raise ValueError(
+                    f"Invalid block_shape {block_shape} for parameter shape "
+                    f"{tens_shape}. Elements of block_shape must divide parameter "
+                    f"shape evenly"
+                )
+        # Compute blocked tensor shape
+        if len(block_shape) > 1 and block_shape[1] > 1:
+            return [
+                tens_shape[0] * tens_shape[1] // (block_shape[0] * block_shape[1]),
+                block_shape[0] * block_shape[1],
+                -1,
+            ]
+        else:
+            return [tens_shape[0] // block_shape[0], block_shape[0], -1]
+
+
+def get_mask_creator_default(mask_type: Union[str, List[int]]) -> PruningMaskCreator:
+    """
+    :param mask_type: type of mask creator to use, can be 'unstructured', for
+        unstructured mask creator, 'block4' for 1x4 block pruning, or a list of two
+        integers for custom block pruning (does not support padding)
+    :return: mask creator object created from the mask type
+    """
+    if mask_type == "unstructured":
+        return UnstructuredPruningMaskCreator()
+    elif mask_type == "block4":
+        return FourBlockMaskCreator()
+    elif isinstance(mask_type, List):
+        if not all(isinstance(val, int) for val in mask_type):
+            raise ValueError(
+                "all values in list specification of BlockMaskCreator must be integers "
+                f"found {mask_type}"
+            )
+        if len(mask_type) != 2:
+            raise ValueError(
+                "expected list of length 2 for specification of BlockMaskCreator, "
+                f"got list with length {len(mask_type)}, mask_type={mask_type}"
+            )
+        return BlockMaskCreator(mask_type)
+    else:
+        raise ValueError(
+            f"Unknown mask_type {mask_type}. Supported mask types include "
+            "'unstructured' and 'block'"
+        )

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

@@ -20,9 +20,8 @@
 
 from sparseml.pytorch.optim.modifier import ModifierProp, PyTorchModifierYAML
 from sparseml.pytorch.sparsification.pruning.mask_creator import (
-    FourBlockMaskCreator,
     PruningMaskCreator,
-    UnstructuredPruningMaskCreator,
+    get_mask_creator_default,
 )
 from sparseml.pytorch.sparsification.pruning.modifier_pruning_base import (
     BasePruningModifier,
@@ -70,9 +69,9 @@ class ACDCPruningModifier(BasePruningModifier):
         immediately after or doing some other prune. Default is True
     :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 sparsity (options: ['unstructured',
-        'channel', 'filter']), List to define block shape of a parameters in and out
-         channels, or a SparsityMaskCreator object. default is 'unstructured'
+    :param mask_type: String to define type of sparsity to apply. May be 'unstructred'
+        for unstructured pruning or 'block4' for four block pruning or a list of two
+        integers for a custom block shape. Default is 'unstructured'
     :param momentum_buffer_reset: set True to reset momentum buffer
         before algorithm enters a consecutive decompression phase.
         According to the paper:
@@ -225,15 +224,7 @@ def _get_mask_creator(
         :param params: list of Parameters to be masked
         :return: mask creator object to be used by this pruning algorithm
         """
-        if self._mask_type == "unstructured":
-            return UnstructuredPruningMaskCreator()
-        elif self._mask_type == "block":
-            return FourBlockMaskCreator()
-        else:
-            raise ValueError(
-                f"Unknown mask_type {self._mask_type}. Supported mask types include "
-                "'unstructured' and 'block'"
-            )
+        return get_mask_creator_default(self.mask_type)
 
     @staticmethod
     def _reset_momentum_buffer(optimizer):

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

@@ -24,9 +24,8 @@
 
 from sparseml.pytorch.optim.modifier import PyTorchModifierYAML
 from sparseml.pytorch.sparsification.pruning.mask_creator import (
-    FourBlockMaskCreator,
     PruningMaskCreator,
-    UnstructuredPruningMaskCreator,
+    get_mask_creator_default,
 )
 from sparseml.pytorch.sparsification.pruning.modifier_pruning_base import (
     BaseGradualPruningModifier,
@@ -104,8 +103,9 @@ class GMPruningModifier(BaseGradualPruningModifier, BaseGMPruningModifier):
         [linear, cubic, inverse_cubic]
     :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 sparsity to apply. May be 'unstructured'
-        for unstructured pruning or 'block' for four block pruning
+    :param mask_type: String to define type of sparsity to apply. May be 'unstructred'
+        for unstructured pruning or 'block4' for four block pruning or a list of two
+        integers for a custom block shape. Default is 'unstructured'
     """
 
     def __init__(
@@ -152,15 +152,7 @@ def _get_mask_creator(
         :param params: list of parameters to be masked
         :return: mask creator object to be used by this pruning algorithm
         """
-        if self.mask_type == "unstructured":
-            return UnstructuredPruningMaskCreator()
-        elif self.mask_type == "block":
-            return FourBlockMaskCreator()
-        else:
-            raise ValueError(
-                f"Unknown mask_type {self.mask_type}. Supported mask types include "
-                "'unstructured' and 'block'"
-            )
+        return get_mask_creator_default(self.mask_type)
 
     def _get_scorer(self, params: List[Parameter]) -> PruningParamsScorer:
         """
@@ -221,7 +213,8 @@ class MagnitudePruningModifier(GMPruningModifier):
     :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 sparsity to apply. May be 'unstructred'
-        for unstructured pruning or 'block' for four block pruning
+        for unstructured pruning or 'block4' for four block pruning or a list of two
+        integers for a custom block shape. Default is 'unstructured'
     """
 
     # just an alias for GMPruningModifier
@@ -274,7 +267,8 @@ class GlobalMagnitudePruningModifier(GMPruningModifier):
     :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 sparsity to apply. May be 'unstructred'
-        for unstructured pruning or 'block' for four block pruning
+        for unstructured pruning or 'block4' for four block pruning or a list of two
+        integers for a custom block shape. Default is 'unstructured'
     """
 
     def __init__(

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

@@ -32,9 +32,8 @@
 import GPUtil
 from sparseml.pytorch.optim.modifier import ModifierProp, PyTorchModifierYAML
 from sparseml.pytorch.sparsification.pruning.mask_creator import (
-    FourBlockMaskCreator,
     PruningMaskCreator,
-    UnstructuredPruningMaskCreator,
+    get_mask_creator_default,
 )
 from sparseml.pytorch.sparsification.pruning.modifier_pruning_base import (
     BaseGradualPruningModifier,
@@ -109,9 +108,9 @@ class MFACPruningModifier(BaseGradualPruningModifier):
         [linear, cubic, inverse_cubic]
     :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 sparsity (options: ['unstructured',
-        'channel', 'filter']), List to define block shape of a parameters in and out
-        channels, or a SparsityMaskCreator object. default is 'unstructured'
+    :param mask_type: String to define type of sparsity to apply. May be 'unstructred'
+        for unstructured pruning or 'block4' for four block pruning or a list of two
+        integers for a custom block shape. Default is 'unstructured'
     :param global_sparsity: set True to enable global pruning. if False, pruning will
         be layer-wise. Default is False
     :param use_gradient_buffering: Optional bool to use gradient buffering instead of
@@ -132,9 +131,9 @@ class MFACPruningModifier(BaseGradualPruningModifier):
         Default is 1
     :param available_devices: list of device names to perform computation on. Default
         is empty
-    :param mask_type: String to define type of sparsity (options: ['unstructured',
-        'block']), List to define block shape of a parameters in and out
-         channels, or a SparsityMaskCreator object. default is 'unstructured'
+    :param mask_type: String to define type of sparsity to apply. May be 'unstructred'
+        for unstructured pruning or 'block4' for four block pruning or a list of two
+        integers for a custom block shape. Default is 'unstructured'
     """
 
     def __init__(
@@ -296,15 +295,7 @@ def _get_mask_creator(
         :param params: list of Parameters to be masked
         :return: mask creator object to be used by this pruning algorithm
         """
-        if self._mask_type == "unstructured":
-            return UnstructuredPruningMaskCreator()
-        elif self._mask_type == "block":
-            return FourBlockMaskCreator()
-        else:
-            raise ValueError(
-                f"Unknown mask_type {self._mask_type}. Supported mask types include "
-                "'unstructured' and 'block'"
-            )
+        return get_mask_creator_default(self.mask_type)
 
     def _get_scorer(self, params: List[Parameter]) -> PruningParamsGradScorer:
         """

tests/sparseml/pytorch/sparsification/pruning/test_mask_creator.py

@@ -16,6 +16,7 @@
 import torch
 
 from sparseml.pytorch.sparsification.pruning import (
+    BlockMaskCreator,
     FourBlockMaskCreator,
     GroupedPruningMaskCreator,
     UnstructuredPruningMaskCreator,
@@ -35,6 +36,7 @@
         ([[64, 513]], FourBlockMaskCreator()),
         ([[64, 512, 3, 3]], FourBlockMaskCreator()),
         ([[63, 513, 3, 3]], FourBlockMaskCreator()),
+        ([[64, 512, 3, 3]], BlockMaskCreator([1, 4])),
     ],
 )
 @pytest.mark.parametrize("sparsity_val", [0.0, 0.4, 0.6, 0.9, 0.99, 1.0])
@@ -47,6 +49,7 @@ def test_sparsity_mask_creator(tensor_shape, mask_creator, sparsity_val):
     [
         ([[64, 64, 3, 3]], UnstructuredPruningMaskCreator()),
         ([[64, 512, 3, 3]], FourBlockMaskCreator()),
+        ([[64, 512, 3, 3]], BlockMaskCreator([1, 4])),
     ],
 )
 @pytest.mark.parametrize("sparsity_val", [0.0, 0.4, 0.6, 0.9, 0.99, 1.0])
@@ -74,6 +77,14 @@ def test_sparsity_mask_creator_cuda(tensor_shape, mask_creator, sparsity_val):
             [i * torch.randn(64, 64, 3, 3) for i in range(1, 6)],
             FourBlockMaskCreator(),
         ),
+        (
+            [torch.randn(128, 128, 3, 3), 3 * torch.randn(64, 512)],
+            BlockMaskCreator([1, 4]),
+        ),
+        (
+            [i * torch.randn(64, 64, 3, 3) for i in range(1, 6)],
+            BlockMaskCreator([1, 4]),
+        ),
     ],
 )
 @pytest.mark.parametrize("sparsity_val", [0.0, 0.4, 0.6, 0.9, 0.99, 1.0])
@@ -116,7 +127,37 @@ def test_global_sparsity_mask_creator(tensors, mask_creator, sparsity_val):
             FourBlockMaskCreator(),
             [0.4, 0.6, 0.8, 0.9, 0.95, 0.99],
         ),
+        (
+            [[128, 128, 3, 3], [64, 512]],
+            BlockMaskCreator([1, 4]),
+            [0.8, 0.9],
+        ),
+        (
+            [[64, 64, 3, 3]] * 6,
+            BlockMaskCreator([1, 4]),
+            [0.4, 0.6, 0.8, 0.9, 0.95, 0.99],
+        ),
     ],
 )
 def test_sparsity_mask_creator_mult_tensor(tensor_shapes, mask_creator, sparsity_val):
     sparsity_mask_creator_test(tensor_shapes, mask_creator, sparsity_val, "cpu")
+
+
+@pytest.mark.parametrize(
+    ("tensors"),
+    [
+        [torch.randn(128, 128), torch.randn(128, 512, 3, 3)],
+        [torch.randn(5, 64, 3, 3)],
+    ],
+)
+@pytest.mark.parametrize("sparsity_val", [0.0, 0.4, 0.6, 0.9, 0.99, 1.0])
+def test_four_block_mask_creator_matches_block(tensors, sparsity_val):
+    mask_creator_1 = FourBlockMaskCreator()
+    mask_creator_2 = BlockMaskCreator([1, 4])
+
+    masks_1 = mask_creator_1.create_sparsity_masks(tensors, sparsity_val)
+    masks_2 = mask_creator_2.create_sparsity_masks(tensors, sparsity_val)
+
+    for mask_1, mask_2 in zip(masks_1, masks_2):
+        assert mask_1.shape == mask_2.shape
+        assert torch.all(mask_1 == mask_2)