Skip to content
This repository was archived by the owner on Jun 3, 2025. It is now read-only.

[cherry-pick] migrate LayerInfo, Result, and ModelResult to pydantic BaseModel (#405) #411

Merged
markurtz merged 1 commit into from
Oct 19, 2021
Merged

[cherry-pick] migrate LayerInfo, Result, and ModelResult to pydantic BaseModel (#405) #411

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
212 changes: 78 additions & 134 deletions src/sparseml/sparsification/model_info.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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(
Expand Down Expand Up @@ -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):
Expand All @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand Down
38 changes: 23 additions & 15 deletions tests/sparseml/sparsification/test_model_info.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
),
(
Expand All @@ -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,
},
),
(
Expand All @@ -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,
},
),
(
Expand All @@ -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
Expand Down Expand Up @@ -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}),
},
),
{
Expand All @@ -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
Expand Down