mct-nightly 2.2.0.20250113.527__py3-none-any.whl → 2.2.0.20250114.84821__py3-none-any.whl

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

@@ -12,29 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, Any
+from typing import Dict, Any, Set
 
 import numpy as np
 
 
 class RUTarget(Enum):
     """
-    Targets for which we define Resource Utilization metrics for mixed-precision search.
-    For each target that we care to consider in a mixed-precision search, there should be defined a set of
-    resource utilization computation function, resource utilization aggregation function,
-    and resource utilization target (within a ResourceUtilization object).
-
-    Whenever adding a resource utilization metric to ResourceUtilization class we should add a matching target to this enum.
-
-    WEIGHTS - Weights memory ResourceUtilization metric.
-
-    ACTIVATION - Activation memory ResourceUtilization metric.
-
-    TOTAL - Total memory ResourceUtilization metric.
-
-    BOPS - Total Bit-Operations ResourceUtilization Metric.
+    Resource Utilization targets for mixed-precision search.
 
+    WEIGHTS - Weights memory.
+    ACTIVATION - Activation memory.
+    TOTAL - Total memory.
+    BOPS - Total Bit-Operations.
     """
 
     WEIGHTS = 'weights'
@@ -43,34 +35,20 @@ class RUTarget(Enum):
     BOPS = 'bops'
 
 
+@dataclass
 class ResourceUtilization:
     """
     Class to represent measurements of performance.
-    """
-
-    def __init__(self,
-                 weights_memory: float = np.inf,
-                 activation_memory: float = np.inf,
-                 total_memory: float = np.inf,
-                 bops: float = np.inf):
-        """
-
-        Args:
-            weights_memory: Memory of a model's weights in bytes. Note that this includes only coefficients that should be quantized (for example, the kernel of Conv2D in Keras will be affected by this value, while the bias will not).
-            activation_memory: Memory of a model's activation in bytes, according to the given activation resource utilization metric.
-            total_memory: The sum of model's activation and weights memory in bytes, according to the given total resource utilization metric.
-            bops: The total bit-operations in the model.
-        """
-        self.weights_memory = weights_memory
-        self.activation_memory = activation_memory
-        self.total_memory = total_memory
-        self.bops = bops
 
-    def __repr__(self):
-        return f"Weights_memory: {self.weights_memory}, " \
-               f"Activation_memory: {self.activation_memory}, " \
-               f"Total_memory: {self.total_memory}, " \
-               f"BOPS: {self.bops}"
+    weights_memory: Memory of a model's weights in bytes.
+    activation_memory: Memory of a model's activation in bytes.
+    total_memory: The sum of model's activation and weights memory in bytes.
+    bops: The total bit-operations in the model.
+    """
+    weights_memory: float = np.inf
+    activation_memory: float = np.inf
+    total_memory: float = np.inf
+    bops: float = np.inf
 
     def weight_restricted(self):
         return self.weights_memory < np.inf
@@ -93,34 +71,30 @@ class ResourceUtilization:
                 RUTarget.TOTAL: self.total_memory,
                 RUTarget.BOPS: self.bops}
 
-    def set_resource_utilization_by_target(self, ru_mapping: Dict[RUTarget, float]):
+    def is_satisfied_by(self, ru: 'ResourceUtilization') -> bool:
         """
-        Setting a ResourceUtilization object values for each ResourceUtilization target in the given dictionary.
+        Checks whether another ResourceUtilization object satisfies the constraints defined by the current object.
 
         Args:
-            ru_mapping: A mapping from a RUTarget to a matching resource utilization value.
+            ru: A ResourceUtilization object to check against the current object.
 
+        Returns:
+            Whether all constraints are satisfied.
         """
-        self.weights_memory = ru_mapping.get(RUTarget.WEIGHTS, np.inf)
-        self.activation_memory = ru_mapping.get(RUTarget.ACTIVATION, np.inf)
-        self.total_memory = ru_mapping.get(RUTarget.TOTAL, np.inf)
-        self.bops = ru_mapping.get(RUTarget.BOPS, np.inf)
+        return bool(ru.weights_memory <= self.weights_memory and \
+                    ru.activation_memory <= self.activation_memory and \
+                    ru.total_memory <= self.total_memory and \
+                    ru.bops <= self.bops)
 
-    def holds_constraints(self, ru: Any) -> bool:
-        """
-        Checks whether the given ResourceUtilization object holds a set of ResourceUtilization constraints defined by
-        the current ResourceUtilization object.
+    def get_restricted_metrics(self) -> Set[RUTarget]:
+        d = self.get_resource_utilization_dict()
+        return {k for k, v in d.items() if v < np.inf}
 
-        Args:
-            ru: A ResourceUtilization object to check if it holds the constraints.
-
-        Returns: True if all the given resource utilization values are not greater than the referenced resource utilization values.
+    def is_any_restricted(self) -> bool:
+        return bool(self.get_restricted_metrics())
 
-        """
-        if not isinstance(ru, ResourceUtilization):
-            return False
-
-        return ru.weights_memory <= self.weights_memory and \
-               ru.activation_memory <= self.activation_memory and \
-               ru.total_memory <= self.total_memory and \
-               ru.bops <= self.bops
+    def __repr__(self):
+        return f"Weights_memory: {self.weights_memory}, " \
+               f"Activation_memory: {self.activation_memory}, " \
+               f"Total_memory: {self.total_memory}, " \
+               f"BOPS: {self.bops}"