v5 refactor: add profiler.py

neurogolf_solver/profiler.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+"""Static profiling for ONNX models."""
+
+import onnx
+from onnx import numpy_helper
+from .constants import BANNED_OPS, GH, GW
+
+try:
+    from neurogolf_utils import score_network as _score_network_official
+    HAS_ONNX_TOOL = True
+except ImportError:
+    HAS_ONNX_TOOL = False
+
+
+def score_network(path):
+    """Score network using official tool or fallback."""
+    if HAS_ONNX_TOOL:
+        try:
+            return _score_network_official(path)
+        except:
+            pass
+    return _static_profile(path)
+
+
+def _static_profile(path):
+    """Static profiling fallback."""
+    try:
+        model = onnx.load(path)
+    except:
+        return None, None, None
+    tensors = {}
+    params = 0
+    nbytes = 0
+    macs = 0
+    for init in model.graph.initializer:
+        a = numpy_helper.to_array(init)
+        tensors[init.name] = a
+        params += a.size
+        nbytes += a.nbytes
+    for nd in model.graph.node:
+        if nd.op_type == 'Constant':
+            for attr in nd.attribute:
+                if attr.t and attr.t.ByteSize() > 0:
+                    try:
+                        a = numpy_helper.to_array(attr.t)
+                        if nd.output:
+                            tensors[nd.output[0]] = a
+                        params += a.size
+                        nbytes += a.nbytes
+                    except:
+                        pass
+        if nd.op_type in BANNED_OPS:
+            return None, None, None
+        if nd.op_type == 'Conv' and len(nd.input) >= 2 and nd.input[1] in tensors:
+            w = tensors[nd.input[1]]
+            if w.ndim == 4:
+                co, ci, kh, kw = w.shape
+                macs += co * ci * kh * kw * GH * GW
+    return int(macs), int(nbytes), int(params)