convert Cut into dataclass, add sorted names signature computed once

Author: irenab

model_compression_toolkit/core/common/graph/memory_graph/cut.py

@@ -12,28 +12,41 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from dataclasses import dataclass, field
+
 from typing import List, Set
 
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.graph.memory_graph.memory_element import MemoryElements
 
 
+@dataclass(frozen=True)
 class Cut:
     """
     A Cut object that contains a set of ordered nodes and their memory elements.
+
+    Args:
+        op_order: A list of the cut's nodes (model layers), ordered by their addition to the cut (first-to-last).
+        op_record: A (unordered) set of the nodes in the cut.
+        mem_elements: MemoryElements object which represents the activation tensors of the cut's nodes.
     """
+    op_order: List[BaseNode]
+    op_record: Set[BaseNode]
+    mem_elements: MemoryElements
 
-    def __init__(self, op_order: List[BaseNode], op_record: Set[BaseNode], mem_elements: MemoryElements):
-        """
-        Args:
-            op_order: A list of the cut's nodes (model layers), ordered by their addition to the cut (first-to-last).
-            op_record: A (unordered) set of the nodes in the cut.
-            mem_elements: MemoryElements object which represents the activation tensors of the cut's nodes.
-        """
+    records_names: Set[str] = field(init=False)
+    _sorted_elements_signature: str = field(init=False, default=None)
+
+    def __post_init__(self):
+        # set frozen attributes
+        object.__setattr__(self, 'records_names', {op.name for op in self.op_record})
 
-        self.op_order = op_order
-        self.op_record = op_record
-        self.mem_elements = mem_elements
+    @property
+    def sorted_elements_signature(self):
+        if self._sorted_elements_signature is None:
+            object.__setattr__(self, '_sorted_elements_signature',
+                               '_'.join(sorted([e.node_name for e in self.mem_elements.elements])))
+        return self._sorted_elements_signature
 
     def memory_size(self) -> float:
         """
@@ -42,15 +55,6 @@ def memory_size(self) -> float:
 
         return self.mem_elements.total_size
 
-    def get_record_names(self) -> Set[str]:
-        """
-        Builds a set of the cut nodes' names.
-
-        Returns: a set with the nodes' names.
-        """
-
-        return {op.name for op in self.op_record}
-
     def __eq__(self, other) -> bool:
         """
         Overrides the class equality method.
@@ -71,7 +75,3 @@ def __hash__(self):
 
     def __repr__(self):
         return f"<Cut: Nodes={[e.node_name for e in self.mem_elements.elements]}, size={self.memory_size()}>"  # pragma: no cover
-
-    def get_sorted_node_names(self):
-        """ Return sorted node names of memory elements. """
-        return sorted([e.node_name for e in self.mem_elements.elements])

model_compression_toolkit/core/common/graph/memory_graph/max_cut_astar.py

@@ -233,7 +233,7 @@ def _get_cut_to_expand(self, open_list: Set[Cut], costs: Dict[Cut, float], route
         ordered_cuts_list = sorted(open_list,
                                    key=lambda c: (self.accumulate(costs[c], self.estimate(c, estimate)),
                                                   max_cut_len - len(routes[c]),
-                                                  ''.join(c.get_sorted_node_names())))
+                                                  c.sorted_elements_signature))
 
         assert len(ordered_cuts_list) > 0
         return ordered_cuts_list[0]
@@ -250,8 +250,7 @@ def clean_memory_for_next_step(self, cut: Cut) -> Cut:
 
         """
 
-        cut_records_names = cut.get_record_names()
-        filtered_memory_elements = set(filter(lambda elm: not all(child.name in cut_records_names for child in
+        filtered_memory_elements = set(filter(lambda elm: not all(child.name in cut.records_names for child in
                                                                   self.memory_graph.activation_tensor_children(elm)),
                                               cut.mem_elements.elements))
 

model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_data.py

@@ -68,7 +68,7 @@ def compute_resource_utilization_data(in_model: Any,
                                                      fw_impl,
                                                      fqc,
                                                      bit_width_config=core_config.bit_width_config,
-                                                     mixed_precision_enable=mixed_precision_enable,
+                                                     mixed_precision_enable=True,
                                                      running_gptq=False)
 
     ru_calculator = ResourceUtilizationCalculator(transformed_graph, fw_impl, fw_info)

tests/keras_tests/function_tests/test_graph_max_cut.py

@@ -151,6 +151,5 @@ def test_graph_max_cut_deterministic_order(self):
         assert len(set(max_cut_sizes)) == 1
         # nodes within each cut can be in different order, and cuts can be in different order inside cuts list,
         # but overall the cuts should be identical between different runs
-        sorted_cuts_solutions = [sorted(cut.get_sorted_node_names() for cut in cuts) for cuts in cuts_solutions]
+        sorted_cuts_solutions = [sorted(cut.sorted_elements_signature for cut in cuts) for cuts in cuts_solutions]
         assert all(cuts == sorted_cuts_solutions[0] for cuts in sorted_cuts_solutions[1:])
-