#!/usr/bin/env python3 """ score_model.py — Full profiled scoring that matches Kaggle exactly. Verified: task084=12.273, task153=13.998 (exact match to Kaggle/LEARNING.md) Usage: python score_model.py --model task084.onnx --task-num 84 --task-data-dir /path/to/task-data NOTE: Models with >~500 nodes may hit ORT profiler event limit on some machines. On Kaggle this limit is higher. For local testing, focus on models <200 nodes. """ import onnx import onnxruntime as ort import numpy as np import json import math import os import sys import argparse def sanitize_model(model): """Exact replica of neurogolf_utils.sanitize_model.""" for node in model.graph.node: node.name = node.output[0] if node.output else '' if node.output and "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] if node.output else '' return model def calculate_memory(model, trace_path): """Exact replica of neurogolf_utils.calculate_memory.""" 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 o in node.output: if o: tensor_names.add(o) 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 tname in tensor_names: item = tensor_map.get(tname) if not item: return None if item.type.HasField("sequence_type"): return None if not item.type.HasField("tensor_type"): continue tt = item.type.tensor_type if not tt.HasField("shape"): return None num_el = 1 for dim in tt.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_el *= dim.dim_value if tname in ['input', 'output']: continue np_dtype = onnx.helper.tensor_dtype_to_np_dtype(tt.elem_type) tensor_memory[tname] = num_el * np.dtype(np_dtype).itemsize tensor_dtypes[tname] = np_dtype 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 o in node.output: if o and o != "output": item = tensor_map.get(o) if item is None or not item.type.HasField("tensor_type"): return None 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 calculate_params(model): """Exact replica of neurogolf_utils.calculate_params.""" 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 _EXCLUDED_OP_TYPES = ["LOOP", "SCAN", "NONZERO", "UNIQUE", "SCRIPT", "FUNCTION", "COMPRESS"] def score_model(model_path, task_num, task_data_dir): """Full profiled scoring. Returns (score, memory, params) or (None, None, None).""" task_file = os.path.join(task_data_dir, f'task{task_num:03d}.json') if not os.path.exists(task_file): return None, None, None with open(task_file) as f: data = json.load(f) model = onnx.load(model_path) for node in model.graph.node: if node.op_type.upper() in _EXCLUDED_OP_TYPES: return None, None, None if "Sequence" in node.op_type: return None, None, None sanitized = sanitize_model(model) if sanitized is None: return None, None, None options = ort.SessionOptions() options.enable_profiling = True options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL options.profile_file_prefix = f"{task_num:03}" try: session = ort.InferenceSession(sanitized.SerializeToString(), options) except Exception as e: return None, None, None all_ex = data.get('train', []) + data.get('test', []) + data.get('arc-gen', []) for ex in all_ex: grid = ex['input'] if max(len(grid), max((len(r) for r in grid), default=0)) > 30: continue inp = np.zeros((1, 10, 30, 30), dtype=np.float32) for r, row in enumerate(grid): for c, v in enumerate(row): if r < 30 and c < 30: inp[0][v][r][c] = 1.0 try: session.run(['output'], {'input': inp}) except: pass trace_path = session.end_profiling() memory = calculate_memory(sanitized, trace_path) params = calculate_params(sanitized) os.remove(trace_path) if memory is None or params is None: return None, None, None if memory < 0 or params < 0: return None, None, None return max(1.0, 25.0 - math.log(max(1.0, memory + params))), memory, params if __name__ == '__main__': parser = argparse.ArgumentParser(description="Full profiled scoring (matches Kaggle)") parser.add_argument('--model', required=True, help='Model path') parser.add_argument('--task-num', type=int, required=True, help='Task number') parser.add_argument('--task-data-dir', required=True, help='Task data directory') args = parser.parse_args() score, mem, par = score_model(args.model, args.task_num, args.task_data_dir) if score: print(f"task{args.task_num:03d}: score={score:.3f} (memory={mem:,}, params={par:,})") else: print(f"task{args.task_num:03d}: SCORING FAILED") sys.exit(1)