Base Modifier classes and refactor for framework impls (#399)

Author: bfineran

src/sparseml/keras/optim/modifier_epoch.py

@@ -17,13 +17,14 @@
 """
 
 from sparseml.keras.optim.modifier import KerasModifierYAML, ScheduledModifier
+from sparseml.sparsification import EpochRangeModifier as BaseEpochRangeModifier
 
 
 __all__ = ["EpochRangeModifier"]
 
 
 @KerasModifierYAML()
-class EpochRangeModifier(ScheduledModifier):
+class EpochRangeModifier(BaseEpochRangeModifier, ScheduledModifier):
     """
     Simple modifier to set the range of epochs to train over for
     the recalibration process.
@@ -45,6 +46,6 @@ def __init__(
         :param start_epoch: The epoch to start the modifier at
         :param end_epoch: The epoch to end the modifier at
         """
-        super().__init__(
+        super(EpochRangeModifier, self).__init__(
             start_epoch=start_epoch, end_epoch=end_epoch, end_comparator=-1
         )

src/sparseml/keras/optim/modifier_lr.py

@@ -26,7 +26,10 @@
     ScheduledUpdateModifier,
 )
 from sparseml.keras.utils import KerasLogger, LoggerSettingCallback, LoggingMode, keras
-from sparseml.optim import LearningRate, SetLearningRate
+from sparseml.sparsification import LearningRateModifier as BaseLearningRateModifier
+from sparseml.sparsification import (
+    SetLearningRateModifier as BaseSetLearningRateModifier,
+)
 from sparseml.utils import ALL_TOKEN
 
 
@@ -186,7 +189,7 @@ def _is_logging_step(self):
 
 
 @KerasModifierYAML()
-class SetLearningRateModifier(ScheduledModifier, SetLearningRate):
+class SetLearningRateModifier(BaseSetLearningRateModifier, ScheduledModifier):
     """
     Modifier to set the learning rate to a specific value at a certain point
     in the training process. Once that point is reached, will update the optimizer's
@@ -212,7 +215,7 @@ def __init__(
         end_epoch: float = -1,
         log_types: Union[str, List[str]] = ALL_TOKEN,
     ):
-        super().__init__(
+        super(SetLearningRateModifier, self).__init__(
             learning_rate=learning_rate,
             log_types=log_types,
             start_epoch=start_epoch,
@@ -316,7 +319,7 @@ def get_config(self):
 
 
 @KerasModifierYAML()
-class LearningRateModifier(ScheduledUpdateModifier, LearningRate):
+class LearningRateModifier(BaseLearningRateModifier, ScheduledUpdateModifier):
     """
     Modifier to set the learning rate to follow specific schedulers
     within a period of epochs.
@@ -358,7 +361,7 @@ def __init__(
         update_frequency: float = -1.0,
         log_types: Union[str, List[str]] = ALL_TOKEN,
     ):
-        super().__init__(
+        super(LearningRateModifier, self).__init__(
             lr_class=lr_class,
             lr_kwargs=lr_kwargs,
             init_lr=init_lr,

src/sparseml/keras/optim/modifier_params.py

@@ -21,14 +21,13 @@
 
 from tensorflow import Tensor
 
-from sparseml.keras.optim.modifier import (
-    KerasModifierYAML,
-    ModifierProp,
-    ScheduledModifier,
-)
+from sparseml.keras.optim.modifier import KerasModifierYAML, ScheduledModifier
 from sparseml.keras.optim.utils import get_layer_name_from_param
 from sparseml.keras.utils import keras
-from sparseml.utils import ALL_TOKEN, convert_to_bool, flatten_iterable
+from sparseml.sparsification import (
+    TrainableParamsModifier as BaseTrainableParamsModifier,
+)
+from sparseml.utils import ALL_TOKEN, flatten_iterable
 
 
 __all__ = ["TrainableParamsModifier"]
@@ -62,7 +61,7 @@ def on_train_batch_end(self, batch, logs=None):
 
 
 @KerasModifierYAML()
-class TrainableParamsModifier(ScheduledModifier):
+class TrainableParamsModifier(BaseTrainableParamsModifier, ScheduledModifier):
     """
     Modifier to control the params for a given list of parameters.
     Applies the trainability over all epochs.
@@ -93,16 +92,15 @@ def __init__(
         end_epoch: float = -1.0,
     ):
         super(TrainableParamsModifier, self).__init__(
-            start_epoch=-1,
-            end_epoch=-1,
+            params=self._validate_params(params),
+            trainable=trainable,
+            params_strict=params_strict,
+            start_epoch=start_epoch,
+            end_epoch=end_epoch,
             end_comparator=-1,
         )
-        self._params = self._validate_params(params)
         self._layer_names = [get_layer_name_from_param(p) for p in self._params]
-        self._trainable = convert_to_bool(trainable)
-        self._params_strict = convert_to_bool(params_strict)
         self._vars_to_trainable_orig = {}
-        self.validate()
 
     def _validate_params(self, params: Union[str, List[Union[int, str]]]):
         if isinstance(params, str):
@@ -118,66 +116,10 @@ def _validate_params(self, params: Union[str, List[Union[int, str]]]):
             )
         )
 
-    @ModifierProp()
-    def params(self) -> Union[str, List[str]]:
-        """
-        :return: 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. Can also use the token __ALL__ to specify all
-            params
-        """
-        return self._params
-
-    @params.setter
-    def params(self, value: Union[str, List[str]]):
-        """
-        :param value: 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. Can also use the token __ALL__ to specify all
-            params
-        """
-        self._params = self._validate_params(value)
-        self.validate()
-
     @property
     def layer_names(self) -> List[str]:
         return self._layer_names
 
-    @ModifierProp()
-    def trainable(self) -> bool:
-        """
-        :return: True if the param(s) should be made trainable,
-            False to make them non-trainable
-        """
-        return self._trainable
-
-    @trainable.setter
-    def trainable(self, value: bool):
-        """
-        :param value: True if the param(s) should be made trainable,
-            False to make them non-trainable
-        """
-        self._trainable = value
-        self.validate()
-
-    @ModifierProp()
-    def params_strict(self) -> bool:
-        """
-        :return: True if the given param(s) must be found in each layer and
-            will raise an err if not found,
-            False if missing params are ok and will not raise an err
-        """
-        return self._params_strict
-
-    @params_strict.setter
-    def params_strict(self, value: bool):
-        """
-        :param value: True if the given param(s) must be found in each layer and
-            will raise an err if not found,
-            False if missing params are ok and will not raise an err
-        """
-        self._params_strict = value
-
     def validate(self):
         """
         Validate the values of the params for the current instance are valid