pytorch manager - group modifiers by stage in apply (#873)

* pytorch manager - group modifiers by stage in apply

* update docstring

* wrap modifier in iterable if necessary

Author: bfineran

src/sparseml/pytorch/optim/manager.py

@@ -19,7 +19,8 @@
 """
 
 import logging
-from typing import Any, Dict, List, Optional, Union
+import math
+from typing import Any, Dict, Iterable, List, Optional, Union
 
 import torch
 from torch import Tensor
@@ -373,6 +374,45 @@ def load_state_dict(self, state_dict: Dict[str, Dict], strict: bool = True):
 
             modifiers_index[key].load_state_dict(val)
 
+    def apply(
+        self,
+        module: Module,
+        epoch: float = math.inf,
+        loggers: Optional[LoggerManager] = None,
+        finalize: bool = True,
+        **kwargs,
+    ):
+        """
+        Applies the lifecycle of each stage in the manager/recipe
+        by calling into initialize and finalize for each modifier for each stage
+
+        :param module: the PyTorch model/module to modify
+        :param epoch: the epoch to apply the modifier at, defaults to math.inf (end)
+        :param loggers: Optional logger manager to log the modification process to
+        :param finalize: True to invoke finalize after initialize, False otherwise.
+            If training after one shot, set finalize=False to keep modifiers applied.
+        :param kwargs: Optional kwargs to support specific arguments
+            for individual modifiers (passed to initialize and finalize).
+        """
+        if not self.initialized:
+            super().initialize(module, epoch, loggers, **kwargs)
+            self._initialize_epoch = epoch
+
+        modifier_lists = (
+            self._modifiers
+            if isinstance(self._modifiers, List)
+            else list(self._modifiers.values())
+        )
+
+        for modifier_list in modifier_lists:
+
+            self._initialize_modifiers(
+                modifier_list, module, epoch, loggers=loggers, **kwargs
+            )
+
+            if finalize:
+                self._finalize_modifiers(modifier_list, module, **kwargs)
+
     def apply_structure(
         self,
         module: Module,
@@ -422,12 +462,9 @@ def initialize(
         super().initialize(module, epoch, loggers, **kwargs)
         self._initialize_epoch = epoch
 
-        for mod in self.iter_modifiers():
-            if mod.initialized:
-                # check in case modifier was initialized from apply_structure
-                continue
-
-            mod.initialize(module, epoch, loggers, **kwargs)
+        self._initialize_modifiers(
+            self.iter_modifiers(), module, epoch, loggers, **kwargs
+        )
 
     def initialize_loggers(self, loggers: Union[None, LoggerManager, List[BaseLogger]]):
         """
@@ -521,8 +558,7 @@ def finalize(
         """
         super().finalize(module, reset_loggers, **kwargs)
 
-        for mod in self.iter_modifiers():
-            mod.finalize(module, reset_loggers, **kwargs)
+        self._finalize_modifiers(self.iter_modifiers(), module, reset_loggers, **kwargs)
 
     def update(
         self,
@@ -635,3 +671,34 @@ def optimizer_post_step(
                 continue
 
             mod.optimizer_post_step(module, optimizer, epoch, steps_per_epoch)
+
+    def _initialize_modifiers(
+        self,
+        modifiers: Iterable[Modifier],
+        module: Module,
+        epoch: float = 0,
+        loggers: Union[None, LoggerManager, List[BaseLogger]] = None,
+        **kwargs,
+    ):
+        if isinstance(modifiers, Modifier):
+            modifiers = [modifiers]
+
+        for mod in modifiers:
+            if mod.initialized:
+                # check in case modifier was initialized from apply_structure
+                continue
+
+            mod.initialize(module, epoch, loggers, **kwargs)
+
+    def _finalize_modifiers(
+        self,
+        modifiers: Iterable[Modifier],
+        module: Optional[Module] = None,
+        reset_loggers: bool = True,
+        **kwargs,
+    ):
+        if isinstance(modifiers, Modifier):
+            modifiers = [modifiers]
+
+        for mod in modifiers:
+            mod.finalize(module, reset_loggers, **kwargs)