coreml-complexity-analyzer 0.1.0__py3-none-any.whl

coreml_complexity_analyzer/flops_analyzer.py

@@ -0,0 +1,550 @@
+# Copyright (c) 2025, Devavarapu Yashwanth, Ireddi Rakshitha. All rights reserved.
+#
+# Use of this source code is governed by a BSD-3-clause license that can be
+# found in the LICENSE file.
+
+"""
+FLOPS Analyzer for Core ML Models.
+
+Provides floating-point operations (FLOPS) estimation for MIL operations
+in Core ML models including convolutions, linear layers, matrix
+multiplications, and element-wise operations.
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple, Any
+import numpy as np
+
+import coremltools as ct
+from coremltools.models import MLModel
+
+
+@dataclass
+class LayerFLOPS:
+    """FLOPS metrics for a single layer/operation."""
+
+    name: str
+    op_type: str
+    flops: int
+    mac_ops: int
+    input_shapes: List[Tuple[int, ...]]
+    output_shapes: List[Tuple[int, ...]]
+    attributes: Dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def gflops(self) -> float:
+        """Return FLOPS in billions."""
+        return self.flops / 1e9
+
+    @property
+    def mflops(self) -> float:
+        """Return FLOPS in millions."""
+        return self.flops / 1e6
+
+
+class FLOPSAnalyzer:
+    """
+    Analyzer for estimating FLOPS of Core ML models.
+
+    Calculates the number of floating-point operations required to execute
+    each operation in a Core ML model. This is useful for:
+
+    - Comparing model efficiency before and after optimization
+    - Estimating inference latency on different hardware
+    - Identifying computational bottlenecks
+
+    Parameters
+    ----------
+    model : MLModel or str
+        A loaded MLModel instance or path to a .mlmodel/.mlpackage file.
+
+    Examples
+    --------
+    >>> import coremltools as ct
+    >>> from coreml_complexity_analyzer import FLOPSAnalyzer
+    >>>
+    >>> model = ct.models.MLModel("my_model.mlpackage")
+    >>> analyzer = FLOPSAnalyzer(model)
+    >>> total_flops = analyzer.get_total_flops()
+    >>> print(f"Total FLOPS: {total_flops:,}")
+    >>>
+    >>> breakdown = analyzer.get_layer_breakdown()
+    >>> for layer in breakdown[:5]:
+    ...     print(f"{layer.name}: {layer.mflops:.2f} MFLOPS")
+    """
+
+    def __init__(self, model):
+        if isinstance(model, str):
+            self._model = MLModel(model)
+        else:
+            self._model = model
+
+        self._spec = self._model.get_spec()
+        self._layer_flops: List[LayerFLOPS] = []
+        self._analyzed = False
+
+    def analyze(self) -> List[LayerFLOPS]:
+        """
+        Analyze the model and compute FLOPS for each operation.
+
+        Returns
+        -------
+        List[LayerFLOPS]
+            List of FLOPS information for each layer.
+        """
+        self._layer_flops = []
+
+        if self._spec.WhichOneof("Type") == "mlProgram":
+            self._analyze_ml_program()
+        elif self._spec.WhichOneof("Type") in (
+            "neuralNetwork",
+            "neuralNetworkClassifier",
+            "neuralNetworkRegressor",
+        ):
+            self._analyze_neural_network()
+        else:
+            raise ValueError(
+                f"Unsupported model type: {self._spec.WhichOneof('Type')}"
+            )
+
+        self._analyzed = True
+        return self._layer_flops
+
+    def get_total_flops(self) -> int:
+        """
+        Get total FLOPS for the entire model.
+
+        Returns
+        -------
+        int
+            Total FLOPS count.
+        """
+        if not self._analyzed:
+            self.analyze()
+        return sum(layer.flops for layer in self._layer_flops)
+
+    def get_total_mac_ops(self) -> int:
+        """
+        Get total multiply-accumulate operations for the entire model.
+
+        Returns
+        -------
+        int
+            Total MAC operations count.
+        """
+        if not self._analyzed:
+            self.analyze()
+        return sum(layer.mac_ops for layer in self._layer_flops)
+
+    def get_layer_breakdown(self) -> List[LayerFLOPS]:
+        """
+        Get per-layer FLOPS breakdown sorted by FLOPS descending.
+
+        Returns
+        -------
+        List[LayerFLOPS]
+            Sorted list of LayerFLOPS objects.
+        """
+        if not self._analyzed:
+            self.analyze()
+        return sorted(self._layer_flops, key=lambda x: x.flops, reverse=True)
+
+    def get_flops_by_op_type(self) -> Dict[str, int]:
+        """
+        Get FLOPS aggregated by operation type.
+
+        Returns
+        -------
+        Dict[str, int]
+            Operation type to total FLOPS mapping.
+        """
+        if not self._analyzed:
+            self.analyze()
+
+        flops_by_type: Dict[str, int] = {}
+        for layer in self._layer_flops:
+            flops_by_type[layer.op_type] = (
+                flops_by_type.get(layer.op_type, 0) + layer.flops
+            )
+
+        return dict(sorted(flops_by_type.items(), key=lambda x: x[1], reverse=True))
+
+    def _analyze_ml_program(self):
+        """Analyze FLOPS for ML Program format models."""
+        program = self._spec.mlProgram
+
+        for func in program.functions.values():
+            for block_name, block in func.block_specializations.items():
+                for op in block.operations:
+                    layer_flops = self._compute_operation_flops(op)
+                    if layer_flops.flops > 0:
+                        self._layer_flops.append(layer_flops)
+
+    def _analyze_neural_network(self):
+        """Analyze FLOPS for Neural Network format models."""
+        nn_spec = None
+        if self._spec.HasField("neuralNetwork"):
+            nn_spec = self._spec.neuralNetwork
+        elif self._spec.HasField("neuralNetworkClassifier"):
+            nn_spec = self._spec.neuralNetworkClassifier
+        elif self._spec.HasField("neuralNetworkRegressor"):
+            nn_spec = self._spec.neuralNetworkRegressor
+
+        if nn_spec is None:
+            return
+
+        for layer in nn_spec.layers:
+            layer_flops = self._compute_nn_layer_flops(layer)
+            if layer_flops.flops > 0:
+                self._layer_flops.append(layer_flops)
+
+    def _compute_operation_flops(self, op) -> LayerFLOPS:
+        """Compute FLOPS for a single MIL operation."""
+        op_type = op.type
+        op_name = op.outputs[0].name if op.outputs else "unknown"
+
+        input_shapes = self._extract_input_shapes(op)
+        output_shapes = self._extract_output_shapes(op)
+
+        flops, mac_ops = self._calculate_flops_for_op_type(
+            op_type, op, input_shapes, output_shapes
+        )
+
+        return LayerFLOPS(
+            name=op_name,
+            op_type=op_type,
+            flops=flops,
+            mac_ops=mac_ops,
+            input_shapes=input_shapes,
+            output_shapes=output_shapes,
+            attributes=self._extract_attributes(op),
+        )
+
+    def _compute_nn_layer_flops(self, layer) -> LayerFLOPS:
+        """Compute FLOPS for a Neural Network layer."""
+        layer_type = layer.WhichOneof("layer")
+        layer_name = layer.name
+
+        input_shapes = []
+        output_shapes = []
+
+        flops, mac_ops = self._calculate_flops_for_nn_layer(layer_type, layer)
+
+        return LayerFLOPS(
+            name=layer_name,
+            op_type=layer_type,
+            flops=flops,
+            mac_ops=mac_ops,
+            input_shapes=input_shapes,
+            output_shapes=output_shapes,
+        )
+
+    def _calculate_flops_for_op_type(
+        self,
+        op_type: str,
+        op,
+        input_shapes: List[Tuple[int, ...]],
+        output_shapes: List[Tuple[int, ...]],
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS and MAC ops based on operation type."""
+
+        dispatch = {
+            "conv": lambda: self._conv_flops(op, input_shapes, output_shapes),
+            "conv_transpose": lambda: self._conv_flops(op, input_shapes, output_shapes),
+            "linear": lambda: self._linear_flops(op, input_shapes, output_shapes),
+            "matmul": lambda: self._matmul_flops(input_shapes, output_shapes),
+            "softmax": lambda: self._softmax_flops(output_shapes),
+            "einsum": lambda: self._einsum_flops(op, input_shapes),
+        }
+
+        if op_type in dispatch:
+            return dispatch[op_type]()
+        elif op_type in ("add", "sub", "mul", "real_div"):
+            return self._elementwise_flops(output_shapes)
+        elif op_type in ("relu", "sigmoid", "tanh", "gelu", "silu", "softplus"):
+            return self._activation_flops(op_type, output_shapes)
+        elif op_type in ("reduce_sum", "reduce_mean", "reduce_max", "reduce_min"):
+            return self._reduce_flops(input_shapes, output_shapes)
+        elif op_type in ("batch_norm", "instance_norm", "layer_norm"):
+            return self._normalization_flops(op_type, output_shapes)
+        elif op_type in ("max_pool", "avg_pool"):
+            return self._pooling_flops(op, output_shapes)
+        elif op_type in ("squeeze", "expand_dims", "reshape", "transpose", "const"):
+            return (0, 0)
+        else:
+            return self._estimate_flops_from_output(output_shapes)
+
+    def _conv_flops(
+        self,
+        op,
+        input_shapes: List[Tuple[int, ...]],
+        output_shapes: List[Tuple[int, ...]],
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS for convolution operations."""
+        if not output_shapes:
+            return (0, 0)
+
+        output_shape = output_shapes[0]
+
+        kernel_size = 1
+        in_channels = 1
+        groups = 1
+
+        for attr_name, attr_val in self._get_op_attributes(op).items():
+            if "kernel" in attr_name.lower() or "weight" in attr_name.lower():
+                if hasattr(attr_val, "__len__"):
+                    kernel_size = int(np.prod(attr_val))
+            if "groups" in attr_name.lower():
+                groups = int(attr_val) if attr_val else 1
+
+        if input_shapes and len(input_shapes[0]) >= 2:
+            in_channels = input_shapes[0][1] if len(input_shapes[0]) > 1 else 1
+
+        output_elements = int(np.prod(output_shape))
+        mac_ops = output_elements * kernel_size * (in_channels // max(groups, 1))
+        flops = 2 * mac_ops
+
+        return (flops, mac_ops)
+
+    def _linear_flops(
+        self,
+        op,
+        input_shapes: List[Tuple[int, ...]],
+        output_shapes: List[Tuple[int, ...]],
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS for linear/fully-connected operations."""
+        if not input_shapes or not output_shapes:
+            return (0, 0)
+
+        batch_size = 1
+        if len(input_shapes[0]) >= 1:
+            batch_size = input_shapes[0][0] if input_shapes[0][0] > 0 else 1
+
+        in_features = input_shapes[0][-1] if input_shapes[0] else 1
+        out_features = output_shapes[0][-1] if output_shapes[0] else 1
+
+        mac_ops = batch_size * in_features * out_features
+        flops = 2 * mac_ops
+
+        return (flops, mac_ops)
+
+    def _matmul_flops(
+        self,
+        input_shapes: List[Tuple[int, ...]],
+        output_shapes: List[Tuple[int, ...]],
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS for matrix multiplication."""
+        if len(input_shapes) < 2:
+            return (0, 0)
+
+        shape_a = input_shapes[0]
+        shape_b = input_shapes[1]
+
+        if not shape_a or not shape_b:
+            return (0, 0)
+
+        m = shape_a[-2] if len(shape_a) >= 2 else 1
+        k = shape_a[-1] if len(shape_a) >= 1 else 1
+        n = shape_b[-1] if len(shape_b) >= 1 else 1
+
+        batch_dims = shape_a[:-2] if len(shape_a) > 2 else ()
+        batch_size = int(np.prod(batch_dims)) if batch_dims else 1
+
+        mac_ops = batch_size * m * n * k
+        flops = 2 * mac_ops
+
+        return (flops, mac_ops)
+
+    def _elementwise_flops(
+        self, output_shapes: List[Tuple[int, ...]]
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS for element-wise operations."""
+        if not output_shapes:
+            return (0, 0)
+
+        total_elements = sum(int(np.prod(s)) for s in output_shapes if s)
+        return (total_elements, 0)
+
+    def _activation_flops(
+        self, op_type: str, output_shapes: List[Tuple[int, ...]]
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS for activation functions."""
+        if not output_shapes:
+            return (0, 0)
+
+        total_elements = sum(int(np.prod(s)) for s in output_shapes if s)
+
+        flops_per_element = {
+            "relu": 1,
+            "sigmoid": 4,
+            "tanh": 4,
+            "gelu": 8,
+            "silu": 5,
+            "softplus": 3,
+        }
+
+        multiplier = flops_per_element.get(op_type, 1)
+        return (total_elements * multiplier, 0)
+
+    def _softmax_flops(
+        self, output_shapes: List[Tuple[int, ...]]
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS for softmax operation."""
+        if not output_shapes:
+            return (0, 0)
+
+        total_elements = sum(int(np.prod(s)) for s in output_shapes if s)
+        flops = 5 * total_elements
+        return (flops, 0)
+
+    def _reduce_flops(
+        self,
+        input_shapes: List[Tuple[int, ...]],
+        output_shapes: List[Tuple[int, ...]],
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS for reduction operations."""
+        if not input_shapes:
+            return (0, 0)
+
+        input_elements = sum(int(np.prod(s)) for s in input_shapes if s)
+        return (input_elements, 0)
+
+    def _normalization_flops(
+        self, op_type: str, output_shapes: List[Tuple[int, ...]]
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS for normalization operations."""
+        if not output_shapes:
+            return (0, 0)
+
+        total_elements = sum(int(np.prod(s)) for s in output_shapes if s)
+        flops = 5 * total_elements
+        return (flops, 0)
+
+    def _pooling_flops(
+        self, op, output_shapes: List[Tuple[int, ...]]
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS for pooling operations."""
+        if not output_shapes:
+            return (0, 0)
+
+        kernel_size = 9
+
+        for attr_name, attr_val in self._get_op_attributes(op).items():
+            if "kernel" in attr_name.lower():
+                if hasattr(attr_val, "__len__"):
+                    kernel_size = int(np.prod(attr_val))
+                else:
+                    kernel_size = int(attr_val) ** 2
+
+        output_elements = sum(int(np.prod(s)) for s in output_shapes if s)
+        flops = output_elements * kernel_size
+        return (flops, 0)
+
+    def _einsum_flops(
+        self, op, input_shapes: List[Tuple[int, ...]]
+    ) -> Tuple[int, int]:
+        """Estimate FLOPS for einsum operations."""
+        if len(input_shapes) < 2:
+            return (0, 0)
+
+        total_elements = 1
+        for shape in input_shapes:
+            if shape:
+                total_elements *= int(np.prod(shape))
+
+        estimated_flops = int(np.sqrt(total_elements))
+        return (estimated_flops, estimated_flops // 2)
+
+    def _estimate_flops_from_output(
+        self, output_shapes: List[Tuple[int, ...]]
+    ) -> Tuple[int, int]:
+        """Fallback FLOPS estimation based on output size."""
+        if not output_shapes:
+            return (0, 0)
+
+        total_elements = sum(int(np.prod(s)) for s in output_shapes if s)
+        return (total_elements, 0)
+
+    def _calculate_flops_for_nn_layer(
+        self, layer_type: str, layer
+    ) -> Tuple[int, int]:
+        """Calculate FLOPS for Neural Network layer types."""
+        if layer_type == "convolution":
+            conv = layer.convolution
+            output_channels = conv.outputChannels
+            kernel_channels = conv.kernelChannels
+            kernel_h = conv.kernelSize[0] if conv.kernelSize else 3
+            kernel_w = conv.kernelSize[1] if len(conv.kernelSize) > 1 else kernel_h
+
+            output_size = 1000
+            mac_ops = (
+                output_size * output_channels * kernel_channels * kernel_h * kernel_w
+            )
+            return (2 * mac_ops, mac_ops)
+
+        elif layer_type == "innerProduct":
+            ip = layer.innerProduct
+            in_features = ip.inputChannels
+            out_features = ip.outputChannels
+            mac_ops = in_features * out_features
+            return (2 * mac_ops, mac_ops)
+
+        elif layer_type in ("batchnorm", "instanceNormalization"):
+            return (1000, 0)
+
+        elif layer_type in ("activation", "softmax"):
+            return (1000, 0)
+
+        return (0, 0)
+
+    def _extract_input_shapes(self, op) -> List[Tuple[int, ...]]:
+        """Extract input shapes from an operation."""
+        shapes = []
+        for inp in op.inputs.values():
+            if inp.arguments:
+                arg = inp.arguments[0]
+                if arg.HasField("type"):
+                    shape = self._type_to_shape(arg.type)
+                    if shape:
+                        shapes.append(shape)
+        return shapes
+
+    def _extract_output_shapes(self, op) -> List[Tuple[int, ...]]:
+        """Extract output shapes from an operation."""
+        shapes = []
+        for out in op.outputs:
+            if out.HasField("type"):
+                shape = self._type_to_shape(out.type)
+                if shape:
+                    shapes.append(shape)
+        return shapes
+
+    def _type_to_shape(self, type_spec) -> Optional[Tuple[int, ...]]:
+        """Convert a type specification to a shape tuple."""
+        if type_spec.HasField("tensorType"):
+            dims = []
+            for dim in type_spec.tensorType.dimensions:
+                if dim.HasField("constant"):
+                    dims.append(dim.constant.size)
+                else:
+                    dims.append(-1)
+            return tuple(dims)
+        return None
+
+    def _extract_attributes(self, op) -> Dict[str, Any]:
+        """Extract attributes from an operation."""
+        attrs = {}
+        for attr_name, attr_val in op.attributes.items():
+            if attr_val.HasField("i"):
+                attrs[attr_name] = attr_val.i
+            elif attr_val.HasField("f"):
+                attrs[attr_name] = attr_val.f
+            elif attr_val.HasField("b"):
+                attrs[attr_name] = attr_val.b
+            elif attr_val.HasField("s"):
+                attrs[attr_name] = attr_val.s
+        return attrs
+
+    def _get_op_attributes(self, op) -> Dict[str, Any]:
+        """Get all attributes from an operation."""
+        return self._extract_attributes(op)