migrate LayerInfo, Result, and ModelResult to pydantic BaseModel (#405) (#411)

* migrate LayerInfo and analysis result classes to pydantic

* bump pydantic min version to 1.5.0 (default_factory support)

* removing boilerplate JSON serialization fns for results and layer info in favor of pydantic built-ins

Author: bfineran

setup.py

@@ -43,7 +43,7 @@
     "pandas>=0.25.0",
     "packaging>=20.0",
     "psutil>=5.0.0",
-    "pydantic>=1.0.0",
+    "pydantic>=1.5.0",
     "requests>=2.0.0",
     "scikit-image>=0.15.0",
     "scipy>=1.0.0",

src/sparseml/sparsification/model_info.py

@@ -24,6 +24,7 @@
 from typing import Any, Dict, List, Optional, Union
 
 import numpy
+from pydantic import BaseModel, Field, root_validator
 
 from sparseml.utils import clean_path, create_parent_dirs
 
@@ -36,49 +37,63 @@
 ]
 
 
-class LayerInfo(object):
+class LayerInfo(BaseModel):
     """
     Class for storing properties about a layer in a model
-
-    :param name: unique name of the layer within its model
-    :param op_type: type of layer, i.e. "conv", "linear"
-    :param params: number of non-bias parameters in the layer. must be
-        included for prunable layers
-    :param bias_params: number of bias parameters in the layer
-    :param prunable: True if the layers non-bias parameters can be pruned.
-        Default is False
-    :param flops: optional number of float operations within the layer
-    :param execution_order: optional execution order of the layer within the
-        model. Default is -1
-    :param attributes: optional dictionary of string attribute names to their
-        values
     """
 
-    def __init__(
-        self,
-        name: str,
-        op_type: str,
-        params: Optional[int] = None,
-        bias_params: Optional[int] = None,
-        prunable: bool = False,
-        flops: Optional[int] = None,
-        execution_order: int = -1,
-        attributes: Optional[Dict[str, Any]] = None,
-    ):
+    name: str = Field(
+        title="name",
+        description="unique name of the layer within its model",
+    )
+    op_type: str = Field(
+        title="op_type",
+        description="type of layer, i.e. 'conv', 'linear'",
+    )
+    params: Optional[int] = Field(
+        title="params",
+        default=None,
+        description=(
+            "number of non-bias parameters in the layer. must be included "
+            "for prunable layers"
+        ),
+    )
+    bias_params: Optional[int] = Field(
+        title="bias_params",
+        default=None,
+        description="number of bias parameters in the layer",
+    )
+    prunable: bool = Field(
+        title="prunable",
+        default=False,
+        description="True if the layers non-bias parameters can be pruned",
+    )
+    flops: Optional[int] = Field(
+        title="flops",
+        default=None,
+        description="number of float operations within the layer",
+    )
+    execution_order: int = Field(
+        title="execution_order",
+        default=-1,
+        description="execution order of the layer within the model",
+    )
+    attributes: Optional[Dict[str, Any]] = Field(
+        title="attributes",
+        default=None,
+        description="dictionary of string attribute names to their values",
+    )
+
+    @root_validator(pre=True)
+    def check_params_if_prunable(_, values):
+        prunable = values.get("prunable")
+        params = values.get("params")
         if prunable and not params:
             raise ValueError(
                 f"Prunable layers must have non 0 number of params given {params} "
-                f"for layer {name} with prunable set to {prunable}"
+                f"for layer {values.get('name')} with prunable set to {prunable}"
             )
-
-        self.name = name
-        self.op_type = op_type
-        self.params = params
-        self.bias_params = bias_params
-        self.prunable = prunable
-        self.flops = flops
-        self.execution_order = execution_order
-        self.attributes = attributes or {}
+        return values
 
     @classmethod
     def linear_layer(
@@ -159,112 +174,41 @@ def conv_layer(
             **kwargs,  # TODO: add FLOPS calculation
         )
 
-    @classmethod
-    def from_dict(cls, dictionary: Dict[str, Any]):
-        """
-        :param dictionary: dict serialized by LyaerInfo.from_dict
-        :return: LayerInfo object created from the given dict
-        """
-        dictionary = deepcopy(dictionary)
-        return cls(**dictionary)
 
-    def to_dict(self) -> Dict[str, Any]:
-        """
-        :return: dict representation of this LayerInfo parameters
-        """
-        props = {
-            "name": self.name,
-            "op_type": self.op_type,
-            "prunable": self.prunable,
-            "execution_order": self.execution_order,
-            "attributes": self.attributes,
-        }
-        if self.params is not None:
-            props["params"] = self.params
-        if self.bias_params is not None:
-            props["bias_params"] = self.bias_params
-        if self.flops is not None:
-            props["flops"] = self.flops
-        return props
-
-
-class Result(object):
+class Result(BaseModel):
     """
     Base class for storing the results of an analysis
-
-    :param value: initial value of the result. Defaults to None
-    :param attributes: dict of attributes of this result. Defaults to empty
     """
 
-    def __init__(self, value: Any = None, attributes: Optional[Dict[str, Any]] = None):
-        self.value = value
-        self.attributes = attributes or {}
-
-    @classmethod
-    def from_dict(cls, dictionary: Dict[str, Any]):
-        """
-        :param dictionary: dict serialized by Result.from_dict
-        :return: Result object created from the given dict
-        """
-        dictionary = deepcopy(dictionary)
-        return cls(**dictionary)
-
-    def to_dict(self) -> Dict[str, Any]:
-        """
-        :return: dict representation of this Result
-        """
-        return {"value": self.value, "attributes": self.attributes}
+    value: Any = Field(
+        title="value",
+        default=None,
+        description="initial value of the result",
+    )
+    attributes: Optional[Dict[str, Any]] = Field(
+        title="attributes",
+        default=None,
+        description="dict of attributes of this result",
+    )
 
 
 class ModelResult(Result):
     """
     Class for storing the results of an analysis for an entire model
-
-    :param analysis_type: name of the type of analysis that was performed
-    :param value: initial value of the result. Defaults to None
-    :param layer_results: dict of layer results to initialize for this model.
-        Defaults to empty dict
-    :param attributes: dict of attributes of this result. Defaults to empty
     """
 
-    def __init__(
-        self,
-        analysis_type: str,
-        value: Any = None,
-        layer_results: Dict[str, Result] = None,
-        attributes: Optional[Dict[str, Any]] = None,
-    ):
-        super().__init__(value=value, attributes=attributes)
-
-        self.analysis_type = analysis_type
-        self.layer_results = layer_results or {}
-
-    @classmethod
-    def from_dict(cls, dictionary: Dict[str, Any]):
-        """
-        :param dictionary: dict serialized by ModelResult.from_dict
-        :return: ModelResult object created from the given dict
-        """
-        dictionary = deepcopy(dictionary)
-        dictionary["layer_results"] = dictionary.get("layer_results", {})
-        dictionary["layer_results"] = {
-            layer_name: Result.from_dict(layer_result)
-            for layer_name, layer_result in dictionary["layer_results"].items()
-        }
-        return cls(**dictionary)
-
-    def to_dict(self) -> Dict[str, Any]:
-        """
-        :return: dict representation of this ModelResult
-        """
-        dictionary = super().to_dict()
-        dictionary["analysis_type"] = self.analysis_type
-        dictionary["layer_results"] = {
-            layer_name: layer_result.to_dict()
-            for layer_name, layer_result in self.layer_results.items()
-        }
-
-        return dictionary
+    analysis_type: str = Field(
+        title="analysis_type",
+        description="name of the type of analysis that was performed",
+    )
+    layer_results: Dict[str, Result] = Field(
+        title="layer_results",
+        default_factory=dict,
+        description=(
+            "dict of layer results to initialize for this analysis. should map "
+            "layer name to Result object"
+        ),
+    )
 
 
 class ModelInfo(ABC):
@@ -289,7 +233,7 @@ def __init__(self, model: Any, metadata: Dict[str, Any]):
     @classmethod
     def from_dict(cls, dictionary: Dict[str, Any]):
         """
-        :param dictionary: dict serialized by ModelInfo.from_dict
+        :param dictionary: dict serialized by `dict(ModelInfo(...))`
         :return: ModelInfo object created from the given dict
         """
         dictionary = deepcopy(dictionary)
@@ -298,15 +242,15 @@ def from_dict(cls, dictionary: Dict[str, Any]):
                 "ModelInfo objects serialized as a dict must include a 'layer_info' key"
             )
         layer_info = {
-            name: LayerInfo.from_dict(info)
+            name: LayerInfo.parse_obj(info)
             for name, info in dictionary["layer_info"].items()
         }
 
         model_info = cls(layer_info, metadata=dictionary.get("metadata", {}))
 
         results = dictionary.get("analysis_results", [])
         for result in results:
-            model_result = ModelResult.from_dict(result)
+            model_result = ModelResult.parse_obj(result)
             model_info.add_analysis_result(model_result)
 
         return model_info
@@ -368,8 +312,8 @@ def to_dict(self) -> Dict[str, Any]:
         """
         :return: dict representation of this ModelResult
         """
-        layer_info = {name: info.to_dict() for name, info in self._layer_info.items()}
-        analysis_results = [result.to_dict() for result in self._analysis_results]
+        layer_info = {name: dict(info) for name, info in self._layer_info.items()}
+        analysis_results = [dict(result) for result in self._analysis_results]
         return {
             "metadata": self.metadata,
             "layer_info": layer_info,

tests/sparseml/sparsification/test_model_info.py

@@ -32,13 +32,16 @@ def _test_layer_info_eq(layer_one, layer_two):
     "layer_info,expected_dict",
     [
         (
-            LayerInfo("layers.1", "TestLayer", attributes={"val": 1}),
+            LayerInfo(name="layers.1", op_type="TestLayer", attributes={"val": 1}),
             {
                 "name": "layers.1",
                 "op_type": "TestLayer",
                 "prunable": False,
                 "execution_order": -1,
                 "attributes": {"val": 1},
+                "flops": None,
+                "bias_params": None,
+                "params": None,
             },
         ),
         (
@@ -51,6 +54,7 @@ def _test_layer_info_eq(layer_one, layer_two):
                 "prunable": True,
                 "execution_order": -1,
                 "attributes": {"in_channels": 64, "out_channels": 128},
+                "flops": None,
             },
         ),
         (
@@ -69,6 +73,8 @@ def _test_layer_info_eq(layer_one, layer_two):
                     "stride": 1,
                     "padding": [0, 0, 0, 0],
                 },
+                "bias_params": None,
+                "flops": None,
             },
         ),
         (
@@ -89,14 +95,16 @@ def _test_layer_info_eq(layer_one, layer_two):
                     "stride": 1,
                     "padding": [0, 0, 0, 0],
                 },
+                "bias_params": None,
+                "flops": None,
             },
         ),
     ],
 )
 def test_layer_info_serialization(layer_info, expected_dict):
-    layer_info_dict = layer_info.to_dict()
-    expected_dict_loaded = LayerInfo.from_dict(expected_dict)
-    layer_info_dict_reloaded = LayerInfo.from_dict(layer_info_dict)
+    layer_info_dict = dict(layer_info)
+    expected_dict_loaded = LayerInfo.parse_obj(expected_dict)
+    layer_info_dict_reloaded = LayerInfo.parse_obj(layer_info_dict)
 
     assert type(expected_dict_loaded) is LayerInfo
     assert type(layer_info_dict_reloaded) is LayerInfo
@@ -128,20 +136,20 @@ def _test_model_result_eq(result_one, result_two):
     "model_result,expected_dict",
     [
         (
-            ModelResult("lr_sensitivity", value={0.1: 100, 0.2: 50}),
+            ModelResult(analysis_type="lr_sensitivity", value={0.1: 100, 0.2: 50}),
             {
                 "analysis_type": "lr_sensitivity",
                 "value": {0.1: 100, 0.2: 50},
                 "layer_results": {},
-                "attributes": {},
+                "attributes": None,
             },
         ),
         (
             ModelResult(
-                "pruning_sensitivity",
+                analysis_type="pruning_sensitivity",
                 layer_results={
-                    "net.1": Result({0.0: 0.25, 0.6: 0.2, 0.8: 0.1}),
-                    "net.2": Result({0.0: 0.2, 0.6: 0.2, 0.8: 0.2}),
+                    "net.1": Result(value={0.0: 0.25, 0.6: 0.2, 0.8: 0.1}),
+                    "net.2": Result(value={0.0: 0.2, 0.6: 0.2, 0.8: 0.2}),
                 },
             ),
             {
@@ -150,22 +158,22 @@ def _test_model_result_eq(result_one, result_two):
                 "layer_results": {
                     "net.1": {
                         "value": {0.0: 0.25, 0.6: 0.2, 0.8: 0.1},
-                        "attributes": {},
+                        "attributes": None,
                     },
                     "net.2": {
                         "value": {0.0: 0.2, 0.6: 0.2, 0.8: 0.2},
-                        "attributes": {},
+                        "attributes": None,
                     },
                 },
-                "attributes": {},
+                "attributes": None,
             },
         ),
     ],
 )
 def test_model_result_serialization(model_result, expected_dict):
-    model_result_dict = model_result.to_dict()
-    expected_dict_loaded = ModelResult.from_dict(expected_dict)
-    model_result_dict_reloaded = ModelResult.from_dict(model_result_dict)
+    model_result_dict = dict(model_result)
+    expected_dict_loaded = ModelResult.parse_obj(expected_dict)
+    model_result_dict_reloaded = ModelResult.parse_obj(model_result_dict)
 
     assert type(expected_dict_loaded) is ModelResult
     assert type(model_result_dict_reloaded) is ModelResult