diff --git a/setup.py b/setup.py index dfcdd2bc4fa..fe18fcf92b8 100644 --- a/setup.py +++ b/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", diff --git a/src/sparseml/sparsification/model_info.py b/src/sparseml/sparsification/model_info.py index 128dff71e6d..bbc0af3d7bd 100644 --- a/src/sparseml/sparsification/model_info.py +++ b/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,7 +242,7 @@ 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() } @@ -306,7 +250,7 @@ def from_dict(cls, dictionary: Dict[str, Any]): 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, diff --git a/tests/sparseml/sparsification/test_model_info.py b/tests/sparseml/sparsification/test_model_info.py index 882921670a5..8aefd5c9680 100644 --- a/tests/sparseml/sparsification/test_model_info.py +++ b/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