PyTorch Global Magnitude Pruning (#142)

* refactor mask creators and mask objects to accept lists of parameters instead of single

* global magnitude option for mask creators

* add global pruning option to GMPruningModifier

* create PruningModifier child classes

Author: bfineran

src/sparseml/pytorch/optim/mask_creator_pruning.py

@@ -40,33 +40,41 @@ class PruningMaskCreator(ABC):
     Subclasses should define all methods for creating masks
     """
 
-    def create_sparsity_mask_from_tensor(self, tensor: Tensor) -> Tensor:
+    def create_sparsity_masks_from_tensor(self, tensors: List[Tensor]) -> List[Tensor]:
         """
-        :param tensor: the tensor to calculate a mask based on its values
-        :return: a mask derived from the values of tensor
+        :param tensors: list of tensors to calculate a masks based on their values
+        :return: list of masks derived from each of the given tensors
         """
-        return torch.ne(tensor, 0.0).type(tensor.type())
+        return [torch.ne(tensor, 0.0).type(tensor.type()) for tensor in tensors]
 
     @abstractmethod
-    def create_sparsity_mask_from_abs_threshold(
-        self, tensor: Tensor, threshold: Union[float, Tensor]
-    ) -> Tensor:
+    def create_sparsity_masks_from_abs_threshold(
+        self, tensors: List[Tensor], threshold: Union[float, Tensor]
+    ) -> List[Tensor]:
         """
-        :param tensor: the tensor to calculate a mask from based on the contained
+        :param tensors: list of tensors to calculate a masks based on their contained
             values
         :param threshold: a threshold to determine cutoff for sparsification
-        :return: a mask derived from the tensor and the threshold
+        :return: list of masks derived from each of the given tensors and the threshold
         """
         raise NotImplementedError()
 
     @abstractmethod
-    def create_sparsity_mask(self, tensor: Tensor, sparsity: float) -> Tensor:
+    def create_sparsity_masks(
+        self, tensors: List[Tensor], sparsity: float, global_sparsity: bool = False
+    ) -> List[Tensor]:
         """
-        :param tensor: the tensor to calculate a mask from based on the contained values
+        :param tensors: list of tensors to calculate a masks based on their contained
+            values
         :param sparsity: the desired sparsity to reach within the mask
             (decimal fraction of zeros)
-        :return: a mask (0.0 for values that are masked, 1.0 for values that are
-            unmasked) calculated from the tens such that the desired number of zeros
+        :param global_sparsity: if True, sparsity masks will be created such that the
+            average sparsity across all given tensors is the target sparsity with the
+            lowest global values masked. If False, each tensor will be masked to the
+            target sparsity ranking values within each individual tensor. Default is
+            False
+        :return: list of masks (0.0 for values that are masked, 1.0 for values that are
+            unmasked) calculated from the tensors such that the desired number of zeros
             matches the sparsity.
         """
         raise NotImplementedError()
@@ -76,54 +84,88 @@ class UnstructuredPruningMaskCreator(PruningMaskCreator):
     """
     Class for creating unstructured sparsity masks.
     Masks will be created using unstructured sparsity by pruning weights ranked
-    by their magnitude.
+    by their magnitude.  Each mask will correspond to the given tensor.
     """
 
-    def create_sparsity_mask_from_abs_threshold(
-        self, tensor: Tensor, threshold: Union[float, Tensor]
-    ) -> Tensor:
+    def create_sparsity_masks_from_abs_threshold(
+        self, tensors: List[Tensor], threshold: Union[float, Tensor]
+    ) -> List[Tensor]:
         """
-        :param tensor: the tensor to calculate a mask from based on the contained values
+        :param tensors: list of tensors to calculate a masks based on their contained
+            values
         :param threshold: a threshold at which to mask abs(values) if they are
             less than it or equal
-        :return: a mask (0.0 for values that are masked, 1.0 for values that are
-            unmasked) calculated from the tens abs(values) <= threshold are masked,
+        :return: list of masks (0.0 for values that are masked, 1.0 for values that are
+            unmasked) calculated from the tensors abs(values) <= threshold are masked,
             all others are unmasked
         """
-        return (torch.abs(tensor) > threshold).type(tensor.type())
+        return [
+            (torch.abs(tensor) > threshold).type(tensor.type()) for tensor in tensors
+        ]
 
-    def create_sparsity_mask(self, tensor: Tensor, sparsity: float) -> Tensor:
+    def create_sparsity_masks(
+        self,
+        tensors: List[Tensor],
+        sparsity: float,
+        global_sparsity: bool = False,
+    ) -> List[Tensor]:
         """
-        :param tensor: the tensor to calculate a mask from based on the contained values
+        :param tensors: list of tensors to calculate a mask from based on their
+            contained values
         :param sparsity: the desired sparsity to reach within the mask
             (decimal fraction of zeros)
-        :return: a mask (0.0 for values that are masked, 1.0 for values that are
-            unmasked) calculated from the tens such that the desired number of zeros
-            matches the sparsity. removes the abs lowest values if there are more zeros
-            in the tens than desired sparsity, then will randomly choose the zeros
-        """
-        threshold = self._abs_threshold_from_sparsity(tensor, sparsity)
-
-        if threshold.numel() < 1:
-            return tensor.new_ones(tensor.shape)
-
-        if threshold.item() > 0.0:
-            return (torch.abs(tensor) > threshold).type(tensor.type())
-
-        # too many zeros so will go over the already given sparsity
-        # and choose which zeros to not keep in mask at random
-        zero_indices = (tensor == 0.0).nonzero()
-        rand_indices = list(range(zero_indices.shape[0]))
-        random.shuffle(rand_indices)
-        num_elem = tensor.numel()
-        num_mask = int(num_elem * sparsity)
-        rand_indices = rand_indices[:num_mask]
-        rand_indices = tensor.new_tensor(rand_indices, dtype=torch.int64)
-        zero_indices = zero_indices[rand_indices, :]
-        mask = tensor.new_ones(tensor.shape).type(tensor.type())
-        mask[zero_indices.split(1, dim=1)] = 0
-
-        return mask.type(tensor.type())
+        :param global_sparsity: if True, sparsity masks will be created such that the
+            average sparsity across all given tensors is the target sparsity with the
+            lowest global values masked. If False, each tensor will be masked to the
+            target sparsity ranking values within each individual tensor. Default is
+            False
+        :return: list of masks (0.0 for values that are masked, 1.0 for values that are
+            unmasked) calculated from the tensors such that the desired number of zeros
+            matches the sparsity.  If there are more zeros than the desired sparsity,
+            zeros will be randomly chosen to match the target sparsity
+        """
+        if global_sparsity:
+            # create tensor to make global mask with
+            original_tensors = tensors
+            tensors = [self._flatten_and_stack_tensors(tensors)]
+        else:
+            original_tensors = None
+
+        masks = []
+
+        for tensor in tensors:
+            threshold = self._abs_threshold_from_sparsity(tensor, sparsity)
+
+            if threshold.numel() < 1:
+                masks.append(tensor.new_ones(tensor.shape))
+                continue
+
+            if threshold.item() > 0.0:
+                masks.append((torch.abs(tensor) > threshold).type(tensor.type()))
+                continue
+
+            # too many zeros so will go over the already given sparsity
+            # and choose which zeros to not keep in mask at random
+            zero_indices = (tensor == 0.0).nonzero()
+            rand_indices = list(range(zero_indices.shape[0]))
+            random.shuffle(rand_indices)
+            num_elem = tensor.numel()
+            num_mask = int(num_elem * sparsity)
+            rand_indices = rand_indices[:num_mask]
+            rand_indices = tensor.new_tensor(rand_indices, dtype=torch.int64)
+            zero_indices = zero_indices[rand_indices, :]
+            mask = tensor.new_ones(tensor.shape).type(tensor.type())
+            mask[zero_indices.split(1, dim=1)] = 0
+
+            masks.append(mask.type(tensor.type()))
+
+        if global_sparsity:
+            # unpack global mask into tensor-masks with the original shapes
+            global_mask = masks[0]
+            masks = self._unstack_flattened_tensors(global_mask, original_tensors)
+            del global_mask
+
+        return masks
 
     def _abs_threshold_from_sparsity(self, tensor: Tensor, sparsity: float) -> Tensor:
         """
@@ -146,6 +188,42 @@ def _abs_threshold_from_sparsity(self, tensor: Tensor, sparsity: float) -> Tenso
 
         return sorted_vals[lookup_index]
 
+    def _flatten_and_stack_tensors(self, tensors: List[Tensor]) -> Tensor:
+        total_elements = sum(tensor.numel() for tensor in tensors)
+
+        global_tensor = (
+            tensors[0].new_zeros(total_elements).detach().requires_grad_(False)
+        )
+
+        curr_element = 0
+        for idx, tensor in enumerate(tensors):
+            global_tensor[
+                curr_element : curr_element + tensor.numel()
+            ] = tensor.reshape(-1)
+            curr_element += tensor.numel()
+
+        return global_tensor
+
+    def _unstack_flattened_tensors(
+        self, stacked_tensor: Tensor, original_tensors: List[Tensor]
+    ) -> List[Tensor]:
+        unstacked_tensors = []
+        global_idx = 0
+        for tensor in original_tensors:
+            # unpack global tensor into masks matching original tensor shapes
+            unstacked_tensor = (
+                tensor.new_empty(tensor.numel()).detach().requires_grad_(False)
+            )
+            unstacked_tensor.copy_(
+                stacked_tensor[global_idx : global_idx + tensor.numel()]
+            ).type(tensor.type())
+            unstacked_tensor = unstacked_tensor.reshape(tensor.shape)
+
+            unstacked_tensors.append(unstacked_tensor)
+            global_idx += tensor.numel()
+
+        return unstacked_tensors
+
     def __str__(self):
         return "unstructured"
 
@@ -213,44 +291,71 @@ def _map_mask_to_tensor(
         """
         raise NotImplementedError()
 
-    def create_sparsity_mask_from_tensor(self, tensor: Tensor) -> Tensor:
+    def create_sparsity_masks_from_tensor(self, tensors: List[Tensor]) -> List[Tensor]:
         """
-        :param tensor: the tensor to calculate a mask based on its values
-        :return: a mask derived from the values of tensor grouped by the group_tensor
+        :param tensors: list of tensors to calculate masks based on their values
+        :return: list of masks derived from the values of the tensors grouped by
+            the group_tensor function.
         """
-        grouped_tensor = self.group_tensor(tensor)
-        grouped_mask = super().create_sparsity_mask_from_tensor(grouped_tensor)
-        return self._map_mask_to_tensor(grouped_mask, tensor.shape)
+        masks = []
+        for tensor in tensors:
+            grouped_tensor = self.group_tensor(tensor)
+            grouped_mask = super().create_sparsity_masks_from_tensor([grouped_tensor])[
+                0
+            ]
+            masks.append(self._map_mask_to_tensor(grouped_mask, tensor.shape))
+        return masks
 
-    def create_sparsity_mask_from_abs_threshold(
-        self, tensor: Tensor, threshold: Union[float, Tensor]
-    ) -> Tensor:
+    def create_sparsity_masks_from_abs_threshold(
+        self, tensors: List[Tensor], threshold: Union[float, Tensor]
+    ) -> List[Tensor]:
         """
-        :param tensor: the tensor to calculate a mask from based on the contained
+        :param tensors: list of tensors to calculate masks from based on their contained
             values
         :param threshold: a threshold of group_tensor values to determine cutoff
             for sparsification
-        :return: a mask derived from the tensor and the grouped threshold
-        """
-        grouped_tensor = self.group_tensor(tensor)
-        grouped_mask = super().create_sparsity_mask_from_abs_threshold(
-            grouped_tensor, threshold
-        )
-        return self._map_mask_to_tensor(grouped_mask, tensor.shape)
-
-    def create_sparsity_mask(self, tensor: Tensor, sparsity: float) -> Tensor:
+        :return: list of masks derived from the tensors and the grouped threshold
+        """
+        masks = []
+        for tensor in tensors:
+            grouped_tensor = self.group_tensor(tensor)
+            grouped_mask = super().create_sparsity_masks_from_abs_threshold(
+                [grouped_tensor], threshold
+            )[0]
+            masks.append(self._map_mask_to_tensor(grouped_mask, tensor.shape))
+        return masks
+
+    def create_sparsity_masks(
+        self,
+        tensors: List[Tensor],
+        sparsity: float,
+        global_sparsity: bool = False,
+    ) -> List[Tensor]:
         """
-        :param tensor: the tensor to calculate a mask from based on the contained values
+        :param tensors: list of tensors to calculate masks from based on their contained
+            values
         :param sparsity: the desired sparsity to reach within the mask
             (decimal fraction of zeros)
-        :return: a mask (0.0 for values that are masked, 1.0 for values that are
-            unmasked) calculated from the tens such that the desired number of zeros
-            matches the sparsity and all values mapped to the same group have the
-            same value.
-        """
-        grouped_tensor = self.group_tensor(tensor)
-        grouped_mask = super().create_sparsity_mask(grouped_tensor, sparsity)
-        return self._map_mask_to_tensor(grouped_mask, tensor.shape)
+        :param global_sparsity: if True, sparsity masks will be created such that the
+            average sparsity across all given tensors is the target sparsity with the
+            lowest global values masked. If False, each tensor will be masked to the
+            target sparsity ranking values within each individual tensor. Default is
+            False
+        :return: list of masks (0.0 for values that are masked, 1.0 for values that are
+            unmasked) calculated from the tensors such that the desired number of zeros
+            matches the sparsity and all values mapped to the same group have the same
+            value
+        """
+        grouped_tensors = [self.group_tensor(tensor) for tensor in tensors]
+        grouped_masks = super().create_sparsity_masks(
+            grouped_tensors, sparsity, global_sparsity
+        )
+        masks = [
+            self._map_mask_to_tensor(grouped_mask, tensor.shape)
+            for grouped_mask, tensor in zip(grouped_masks, tensors)
+        ]
+
+        return masks
 
 
 class DimensionSparsityMaskCreator(GroupedPruningMaskCreator):
@@ -286,6 +391,31 @@ def __init__(
                 )
         self._dim = dim  # List[int]
 
+    def create_sparsity_masks(
+        self,
+        tensors: List[Tensor],
+        sparsity: float,
+        global_sparsity: bool = False,
+    ) -> List[Tensor]:
+        """
+        :param tensors: list of tensors to calculate masks from based on their contained
+            values
+        :param sparsity: the desired sparsity to reach within the mask
+            (decimal fraction of zeros)
+        :param global_sparsity: do not set True, unsupported for
+            DimensionSparsityMaskCreator
+        :return: list of masks (0.0 for values that are masked, 1.0 for values that are
+            unmasked) calculated from the tensors such that the desired number of zeros
+            matches the sparsity and all values mapped to the same group have the same
+            value
+        """
+        if global_sparsity:
+            # global sparsity unsupported because channel dims may vary across layers
+            raise ValueError(
+                "global_sparsity not supported for DimensionSparsityMaskCreator"
+            )
+        return super().create_sparsity_masks(tensors, sparsity, global_sparsity=False)
+
     def group_tensor(self, tensor: Tensor) -> Tensor:
         """
         :param tensor: The tensor to transform