Update neurogolf_utils.py to May 14 2026 version (MACs removed from scoring, ORT profiler memory, sanitize_model)

own-solver/neurogolf_utils.py

@@ -12,7 +12,68 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Module containing utilities for the IJCAI-ECAI 2026 NeuroGolf Challenge."""
+"""
+Module containing utilities for the IJCAI-ECAI 2026 NeuroGolf Championship.
+
+Version History:
+* 2026-05-14:
+    * Puts back the _EXCLUDED_OP_TYPES check (thank you @cdeotte and @pavelsavchenkov!)
+    * Applies stronger sanitization to tensor & node names (thank you @linkinpony!)
+    * Rejects duplicate graph.value_info entries with the same tensor name (thank you @pavelsavchenkov!)
+* 2026-05-06:
+    * Scalar parameters are now penalized with unit cost.
+    * Each tensor's memory footprint is set to the maximum size across all runs.
+    * Duplicate node names no longer create parameter undercount.
+    * Tensor names containing ONNX's special "kernel_time" string are disallowed.
+    * Runtime trace file prefixes are specified to prevent profile clobbering.
+    * Multi-input / multi-output graphs disallowed.
+* 2026-05-04:
+    * Sequences and nonpositive tensor dimensions are disallowed.
+    * Accurate shape information derived from the ONNX Runtime Profiler.
+    * MACs no longer contribute to the objective criterion.
+* 2026-05-04:
+    * Sequences and nonpositive tensor dimensions are disallowed.
+    * Accurate shape information derived from the ONNX Runtime Profiler.
+    * MACs no longer contribute to the objective criterion.
+* 2026-04-30:
+    * Compress operators have been banned.
+    * Name collision between tensors and initializers are disallowed.
+    * Functions / custom domains / subgraphs are disallowed.
+    * Zero-cost networks now yield a full 25 points.
+* 2026-04-28:
+    * Constant folding enabled to address the undercounting of parameters.
+    * Our "statically-defined shapes" constaint is now strictly enforced.
+    * Memory footprint calculation is now a sum of static shape sizes.
+    * Nodes with negative parameter counts or MACs are disallowed.
+* 2026-04-21:
+    * Tests with grids larger than 30x30 are ignored.
+    * Nodes with negative memory values are disallowed.
+* 2026-04-15:
+    * Initial version.
+
+Contributors from the Kaggle Community:
+* @anglolodorf
+* @arc144
+* @asalhi
+* @calibrator
+* @cdeotte
+* @hengck23
+* @jazivxt
+* @jiweiliu
+* @kameronkilchrist
+* @kevinyuluo
+* @kosirowada
+* @linkinpony
+* @maxjeblick
+* @mukundan314
+* @pavelsavchenkov
+* @prokaj
+* @robga
+* @shinh0
+* @tonylica
+* @yeoyunsianggeremie
+* @yiheng
+"""
 
 import itertools
 import json
@@ -48,7 +109,7 @@ _COLORS = [
     (146, 117, 86)
 ]
 _DATA_TYPE = onnx.TensorProto.FLOAT
-_EXCLUDED_OP_TYPES = ["LOOP", "SCAN", "NONZERO", "UNIQUE", "SCRIPT", "FUNCTION"]
+_EXCLUDED_OP_TYPES = ["LOOP", "SCAN", "NONZERO", "UNIQUE", "SCRIPT", "FUNCTION", "COMPRESS"]
 _FILESIZE_LIMIT_IN_BYTES = 1.44 * 1024 * 1024
 _GRID_SHAPE = [_BATCH_SIZE, _CHANNELS, _HEIGHT, _WIDTH]
 _IR_VERSION, _OPSET_IMPORTS = 10, [onnx.helper.make_opsetid("", 10)]
@@ -134,6 +195,81 @@ _TASK_ZERO = {
 }
 
 
+def calculate_memory(model, trace_path):
+    onnx.checker.check_model(model, full_check=True)
+    graph = onnx.shape_inference.infer_shapes(model, strict_mode=True).graph
+    if len(graph.input) > 1 or len(graph.output) > 1: return None
+    init_names = {init.name for init in graph.initializer}
+    init_names.update(init.name for init in graph.sparse_initializer)
+    io_names = {t.name for t in list(graph.input) + list(graph.output)}
+    if io_names.intersection(init_names): return None
+    if model.functions: return None
+    for opset in model.opset_import:
+        if opset.domain not in {"", "ai.onnx"}: return None
+    node_outputs = {}
+    tensor_names = set()
+    for node in graph.node:
+        for attr in node.attribute:
+            if attr.type in [onnx.AttributeProto.GRAPH,
+                             onnx.AttributeProto.GRAPHS]:
+                return None
+        node_outputs[node.name] = list(node.output)
+        for output_name in node.output:
+            if output_name: tensor_names.add(output_name)
+    tensor_memory = {}
+    tensor_dtypes = {}
+    tensor_map = {
+        t.name: t for t in list(graph.input) + list(graph.value_info) + list(graph.output)
+    }
+    tensor_names.update(tensor_map.keys())
+    for tensor_name in tensor_names:
+        item = tensor_map.get(tensor_name)
+        if not item: return None
+        if item.type.HasField("sequence_type"): return None
+        if not item.type.HasField("tensor_type"): continue
+        tensor_type = item.type.tensor_type
+        if not tensor_type.HasField("shape"): return None
+        num_elements = 1
+        for dim in tensor_type.shape.dim:
+            if dim.HasField("dim_param"): return None
+            if not dim.HasField("dim_value"): return None
+            if dim.dim_value <= 0: return None
+            num_elements *= dim.dim_value
+        if tensor_name in ['input', 'output']: continue
+        np_dtype = onnx.helper.tensor_dtype_to_np_dtype(tensor_type.elem_type)
+        tensor_memory[tensor_name] = num_elements * np.dtype(np_dtype).itemsize
+        tensor_dtypes[tensor_name] = np_dtype
+
+    # Defensive check to verify uniqueness.
+    seen = set()
+    for item in list(graph.input) + list(graph.value_info) + list(graph.output):
+        if item.name in seen: return None
+        seen.add(item.name)
+    for node in graph.node:
+        for output_name in node.output:
+            if output_name and output_name != "output":
+                item = tensor_map.get(output_name)
+                if item is None or not item.type.HasField("tensor_type"):
+                    return None
+
+    # Retrieve actual tensor shapes via the ONNX Runtime Profiler's JSON Trace.
+    with open(trace_path, 'r') as f:
+        trace_data = json.load(f)
+    for event in trace_data:
+        if event.get("cat") != "Node" or "args" not in event: continue
+        if "output_type_shape" not in event["args"]: continue
+        node_name = event.get("name").replace("_kernel_time", "")
+        if node_name not in node_outputs: continue
+        for i, shape_dict in enumerate(event["args"]["output_type_shape"]):
+            if i >= len(node_outputs[node_name]): continue
+            output_name = node_outputs[node_name][i]
+            if output_name not in tensor_dtypes: continue
+            itemsize = np.dtype(tensor_dtypes[output_name]).itemsize
+            mem = itemsize * sum(math.prod(dims) for dims in shape_dict.values())
+            tensor_memory[output_name] = max(tensor_memory[output_name], mem)
+    return sum(tensor_memory.values())
+
+
 def check_network(filename):
   file_path = pathlib.Path(filename)
   if not file_path.is_file():
@@ -174,23 +310,78 @@ def convert_from_numpy(benchmark):
   return example
 
 
-def score_network(m):
-  model = onnx_tool.loadmodel(m, {'verbose': False})
-  g = model.graph
-  g.graph_reorder_nodes()
-  g.shape_infer(None)
-  g.profile()
-  if not g.valid_profile:
-    print("Error: Invalid profile.")
-    return None, None, None
-  for key in g.nodemap.keys():
-    if g.nodemap[key].op_type.upper() in _EXCLUDED_OP_TYPES:
-      print(f"Error: Op type {g.nodemap[key].op_type} is not permitted.")
-      return None, None, None
-    if g.nodemap[key].memory < 0:
-      print(f"Error: Negative memory value detected.")
-      return None, None, None
-  return int(sum(g.macs)), int(g.memory), int(g.params)
+def calculate_params(model):
+    params = 0
+    for init in model.graph.initializer:
+        if any(d <= 0 for d in init.dims): return None
+        params += math.prod(init.dims)
+    for sparse_init in model.graph.sparse_initializer:
+        if any(d <= 0 for d in sparse_init.values.dims): return None
+        params += math.prod(sparse_init.values.dims)
+    for node in model.graph.node:
+        if node.op_type != 'Constant': continue
+        for attr in node.attribute:
+            if attr.name == 'value':
+                if any(d <= 0 for d in attr.t.dims): return None
+                params += math.prod(attr.t.dims)
+            elif attr.name == 'sparse_value':
+                if any(d <= 0 for d in attr.sparse_tensor.values.dims): return None
+                params += math.prod(attr.sparse_tensor.values.dims)
+            elif attr.name == 'value_floats':
+                params += len(attr.floats)
+            elif attr.name == 'value_ints':
+                params += len(attr.ints)
+            elif attr.name == 'value_strings':
+                params += len(attr.strings)
+    return params
+
+
+def score_network(sanitized, trace_path):
+    for node in sanitized.graph.node:
+        if node.op_type.upper() in _EXCLUDED_OP_TYPES:
+            print(f"Error: Op type {node.op_type} is not permitted.")
+            return None, None
+        if "Sequence" in node.op_type:
+            print(f"Error: Op type {node.op_type} is not permitted.")
+            return None, None
+    return calculate_memory(sanitized, trace_path), calculate_params(sanitized)
+
+
+def sanitize_model(model):
+    for node in model.graph.node:
+        node.name = node.output[0]
+        if "kernel_time" in node.output[0]: return None
+
+    name_map, counter = {}, 0
+
+    def get_safe_name(old_name):
+        nonlocal counter
+        if not old_name or old_name in ["input", "output"]: return old_name
+        if old_name not in name_map:
+            name_map[old_name] = f"safe_name_{counter}"
+            counter += 1
+        return name_map[old_name]
+
+    for inp in model.graph.input:
+        inp.name = get_safe_name(inp.name)
+    for init in model.graph.initializer:
+        init.name = get_safe_name(init.name)
+
+    for node in model.graph.node:
+        for i in range(len(node.input)):
+            node.input[i] = get_safe_name(node.input[i])
+        for i in range(len(node.output)):
+            node.output[i] = get_safe_name(node.output[i])
+        if len(node.output) > 0 and node.output[0]:
+            node.name = node.output[0]
+
+    for out in model.graph.output:
+        out.name = get_safe_name(out.name)
+    for vi in model.graph.value_info:
+        vi.name = get_safe_name(vi.name)
+    for node in model.graph.node:
+        node.name = node.output[0]
+    return model
 
 
 def load_examples(task_num):
@@ -298,7 +489,14 @@ def verify_network(network, task_num, examples):
   onnx.save(network, filename)
   if not check_network(filename): return
   try:
-    session = onnxruntime.InferenceSession(filename)
+    # Load the model, sanitize node names, and enable profiling.
+    sanitized = sanitize_model(onnx.load(filename))
+    if not sanitized: return
+    options = onnxruntime.SessionOptions()
+    options.enable_profiling = True
+    options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
+    options.profile_file_prefix = f"{task_num:03}"
+    session = onnxruntime.InferenceSession(sanitized.SerializeToString(), options)
   except onnxruntime.ONNXRuntimeError as e:
     print(f"Error: Unable to load ONNX model: {e}")
     return
@@ -307,17 +505,19 @@ def verify_network(network, task_num, examples):
   print(f"Results on ARC-AGI examples: {arc_agi_right} pass, {arc_agi_wrong} fail")
   print(f"Results on ARC-GEN examples: {arc_gen_right} pass, {arc_gen_wrong} fail")
   print()
-  macs, memory, params = score_network(filename)
-  if macs is None or memory is None or params is None:
+  memory, params = score_network(sanitized, session.end_profiling())
+  if memory is None or params is None:
+    print("Error: Your network performance could not be measured")
+  if memory < 0 or params < 0:
     print("Error: Your network performance could not be measured")
   elif arc_agi_wrong + arc_gen_wrong == 0:
     print("Your network IS READY for submission!")
     print()
-    print("Performance stats:")
+    print("Performance stats (memory values reported here are approximate):")
     onnx_tool.model_profile(filename)
-    points = max(1.0, 25.0 - math.log(macs + memory + params))
+    points = max(1.0, 25.0 - math.log(max(1.0, memory + params)))
     print()
-    print(f"It appears to require {macs} MACs + {memory} bytes + {params} params, yielding {points:.3f} points.")
+    print(f"It appears to require {memory} bytes + {params} params, yielding {points:.3f} points.")
     print()
     print("Next steps:")
     print(f" * Click the link below to download {filename} onto your local machine.")