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,36 @@
 # 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.
-        """
+    _sorted_elements_signature: str = field(init=False, default=None)
 
-        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:
         """
@@ -71,7 +79,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]

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:])
-