Add optimize_submission.py — unified 4-stage optimizer for new scoring formula

own-solver/optimize_submission.py

@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+"""
+Unified ONNX Optimizer for NeuroGolf — All 4 Stages
+
+Processes submission zip and applies optimizations to maximize score
+under the new formula: score = max(1.0, 25.0 - ln(memory + params))
+
+Stage 1: Rebuild trivial tasks (identity, transpose, color_map, flips)
+Stage 2: Direct Conv→output (eliminate ArgMax+OneHot chain)
+Stage 3: fp16 all weights (halve weight memory contribution)
+Stage 4: Node reduction (remove unused inits, strip metadata)
+
+Usage:
+  python optimize_submission.py --input_zip submission-6043.zip --data_dir ./tasks --output_zip submission_optimized.zip
+"""
+
+import json
+import math
+import os
+import zipfile
+import time
+
+import numpy as np
+import onnx
+import onnxruntime as ort
+from onnx import helper, TensorProto, numpy_helper
+
+GRID_SHAPE = [1, 10, 30, 30]
+DT = TensorProto.FLOAT
+IR = 8
+
+
+def make_model(nodes, inits=None, opset=17):
+    x = helper.make_tensor_value_info("input", DT, GRID_SHAPE)
+    y = helper.make_tensor_value_info("output", DT, GRID_SHAPE)
+    g = helper.make_graph(nodes, "g", [x], [y], initializer=inits or [])
+    return helper.make_model(g, ir_version=IR, opset_imports=[helper.make_opsetid("", opset)])
+
+
+def encode_grid(grid):
+    arr = np.array(grid, dtype=np.int32)
+    h, w = arr.shape
+    t = np.zeros((1, 10, 30, 30), dtype=np.float32)
+    for r in range(h):
+        for c in range(w):
+            v = int(arr[r, c])
+            if 0 <= v < 10:
+                t[0, v, r, c] = 1.0
+    return t
+
+
+def validate_model(model_bytes, examples, max_examples=None):
+    try:
+        opts = ort.SessionOptions()
+        opts.log_severity_level = 3
+        sess = ort.InferenceSession(model_bytes, sess_options=opts, providers=['CPUExecutionProvider'])
+    except Exception:
+        return False
+    exs = examples[:max_examples] if max_examples else examples
+    for ex in exs:
+        try:
+            inp = encode_grid(ex['input'])
+            out = sess.run(['output'], {'input': inp})[0]
+            expected = encode_grid(ex['output'])
+            if not np.array_equal((out > 0.0).astype(np.float32), expected):
+                return False
+        except Exception:
+            return False
+    return True
+
+
+# ═══ STAGE 1 ═══
+
+def stage1_optimize(model_bytes, examples):
+    # Identity
+    m = make_model([helper.make_node('Identity', ['input'], ['output'])])
+    b = m.SerializeToString()
+    if validate_model(b, examples):
+        return b, "S1:identity", 25.0
+
+    # Transpose
+    m = make_model([helper.make_node('Transpose', ['input'], ['output'], perm=[0, 1, 3, 2])])
+    b = m.SerializeToString()
+    if validate_model(b, examples):
+        return b, "S1:transpose", 25.0
+
+    # Flips
+    for axis, name in [(3, 'flip_lr'), (2, 'flip_ud')]:
+        inits = [
+            numpy_helper.from_array(np.array([29], dtype=np.int64), 'st'),
+            numpy_helper.from_array(np.array([np.iinfo(np.int64).min], dtype=np.int64), 'en'),
+            numpy_helper.from_array(np.array([axis], dtype=np.int64), 'ax'),
+            numpy_helper.from_array(np.array([-1], dtype=np.int64), 'sp'),
+        ]
+        m = make_model([helper.make_node('Slice', ['input', 'st', 'en', 'ax', 'sp'], ['output'])], inits)
+        b = m.SerializeToString()
+        if validate_model(b, examples):
+            return b, f"S1:{name}", 21.4
+
+    # Color map detection
+    cm = {}
+    is_cm = True
+    for ex in examples[:10]:
+        inp, out = np.array(ex['input']), np.array(ex['output'])
+        if inp.shape != out.shape:
+            is_cm = False; break
+        for iv, ov in zip(inp.flat, out.flat):
+            iv, ov = int(iv), int(ov)
+            if iv in cm and cm[iv] != ov:
+                is_cm = False; break
+            cm[iv] = ov
+        if not is_cm: break
+
+    if is_cm and cm:
+        is_perm = set(cm.keys()) <= set(range(10)) and set(cm.values()) <= set(range(10))
+        if is_perm:
+            gather_ch = list(range(10))
+            for src, dst in cm.items():
+                if 0 <= src < 10 and 0 <= dst < 10:
+                    gather_ch[dst] = src
+            gi = np.array(gather_ch, dtype=np.int32)
+            inits = [numpy_helper.from_array(gi, 'gi')]
+            m = make_model([helper.make_node('Gather', ['input', 'gi'], ['output'], axis=1)], inits)
+            b = m.SerializeToString()
+            if validate_model(b, examples):
+                return b, "S1:color_perm", 21.1
+
+        W = np.zeros((10, 10, 1, 1), dtype=np.float32)
+        for ic in range(10):
+            oc = cm.get(ic, ic)
+            if 0 <= oc < 10: W[oc, ic, 0, 0] = 1.0
+        inits = [numpy_helper.from_array(W, 'W')]
+        m = make_model([helper.make_node('Conv', ['input', 'W'], ['output'], kernel_shape=[1, 1])], inits)
+        b = m.SerializeToString()
+        if validate_model(b, examples):
+            return b, "S1:color_conv1x1", 18.8
+
+    return None
+
+
+# ═══ STAGE 2 ═══
+
+def stage2_optimize(model_bytes, examples):
+    try:
+        model = onnx.load_from_string(model_bytes)
+    except:
+        return None
+
+    for node in model.graph.node:
+        if node.op_type != 'Conv' or len(node.input) < 2:
+            continue
+        W = None
+        for init in model.graph.initializer:
+            if init.name == node.input[1]:
+                W = numpy_helper.to_array(init); break
+        if W is None or W.ndim != 4 or W.shape[0] != 10 or W.shape[1] != 10:
+            continue
+        if W.shape[2] != W.shape[3]:
+            continue
+        ks = W.shape[2]
+        pad_k = ks // 2
+
+        # Direct conv→output
+        inits = [numpy_helper.from_array(W, 'W')]
+        nodes = [helper.make_node('Conv', ['input', 'W'], ['output'], kernel_shape=[ks, ks], pads=[pad_k]*4)]
+        m = make_model(nodes, inits)
+        b = m.SerializeToString()
+        if validate_model(b, examples):
+            cost = W.size * 4 + W.size
+            return b, f"S2:direct_conv_ks{ks}", max(1.0, 25.0 - math.log(max(1, cost)))
+
+        # With mask (conv_var pattern)
+        inits2 = [numpy_helper.from_array(W, 'W'), numpy_helper.from_array(np.array([1], dtype=np.int64), 'ax')]
+        nodes2 = [
+            helper.make_node('ReduceSum', ['input', 'ax'], ['mask'], keepdims=1),
+            helper.make_node('Conv', ['input', 'W'], ['co'], kernel_shape=[ks, ks], pads=[pad_k]*4),
+            helper.make_node('Mul', ['co', 'mask'], ['output']),
+        ]
+        m2 = make_model(nodes2, inits2)
+        b2 = m2.SerializeToString()
+        if validate_model(b2, examples):
+            cost = 3600 + 36000 + W.size * 4 + 8 + W.size + 1
+            return b2, f"S2:direct_conv_var_ks{ks}", max(1.0, 25.0 - math.log(max(1, cost)))
+
+    return None
+
+
+# ═══ STAGE 3 ═══
+
+def stage3_optimize(model_bytes, examples):
+    try:
+        model = onnx.load_from_string(model_bytes)
+    except:
+        return None
+
+    f32_bytes = sum(numpy_helper.to_array(i).nbytes for i in model.graph.initializer
+                    if numpy_helper.to_array(i).dtype == np.float32 and numpy_helper.to_array(i).size > 10)
+    if f32_bytes < 200:
+        return None
+
+    new_model = onnx.load_from_string(model_bytes)
+    for i, init in enumerate(new_model.graph.initializer):
+        arr = numpy_helper.to_array(init)
+        if arr.dtype == np.float32 and arr.size > 10:
+            new_model.graph.initializer[i].CopyFrom(numpy_helper.from_array(arr.astype(np.float16), name=init.name))
+
+    b = new_model.SerializeToString()
+    if validate_model(b, examples[:15]):
+        if validate_model(b, examples):
+            return b, f"S3:fp16_weights(-{f32_bytes//2//1024}KB)", None
+    return None
+
+
+# ═══ STAGE 4 ═══
+
+def stage4_optimize(model_bytes, examples):
+    try:
+        model = onnx.load_from_string(model_bytes)
+    except:
+        return None
+
+    changed = False
+    used = set()
+    for node in model.graph.node:
+        for inp in node.input: used.add(inp)
+
+    orig = len(model.graph.initializer)
+    new_inits = [i for i in model.graph.initializer if i.name in used]
+    if len(new_inits) < orig:
+        del model.graph.initializer[:]
+        model.graph.initializer.extend(new_inits)
+        changed = True
+
+    if model.doc_string: model.doc_string = ""; changed = True
+    if model.graph.doc_string: model.graph.doc_string = ""; changed = True
+
+    if not changed: return None
+    b = model.SerializeToString()
+    if validate_model(b, examples[:10]):
+        saved = len(model_bytes) - len(b)
+        if saved > 10:
+            return b, f"S4:cleanup(-{saved}B)", None
+    return None
+
+
+# ═══ MAIN ═══
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input_zip', required=True)
+    parser.add_argument('--data_dir', required=True)
+    parser.add_argument('--output_zip', required=True)
+    parser.add_argument('--stages', default='1,2,3,4')
+    args = parser.parse_args()
+    stages = [int(s) for s in args.stages.split(',')]
+
+    models = {}
+    with zipfile.ZipFile(args.input_zip, 'r') as zf:
+        for tid in range(1, 401):
+            fname = f'task{tid:03d}.onnx'
+            if fname in zf.namelist(): models[tid] = zf.read(fname)
+    print(f"Loaded {len(models)} models. Stages: {stages}")
+
+    results = {}
+    counts = {1:0, 2:0, 3:0, 4:0}
+    t0 = time.time()
+
+    for tid in sorted(models.keys()):
+        task_path = os.path.join(args.data_dir, f'task{tid:03d}.json')
+        if not os.path.exists(task_path): continue
+        with open(task_path) as f: task_data = json.load(f)
+        examples = task_data.get('train', []) + task_data.get('test', []) + task_data.get('arc-gen', [])[:30]
+        if not examples: continue
+
+        best = None
+        if 1 in stages:
+            r = stage1_optimize(models[tid], examples)
+            if r: best = r; counts[1] += 1
+        if best is None and 2 in stages:
+            r = stage2_optimize(models[tid], examples)
+            if r: best = r; counts[2] += 1
+        if best is None and 3 in stages:
+            r = stage3_optimize(models[tid], examples)
+            if r: best = r; counts[3] += 1
+        if 4 in stages:
+            target = best[0] if best else models[tid]
+            r = stage4_optimize(target, examples)
+            if r:
+                if best: best = (r[0], best[1]+"+"+r[1], best[2])
+                else: best = r
+                counts[4] += 1
+
+        if best:
+            results[tid] = best[0]
+            score_s = f"score={best[2]:.1f}" if best[2] else ""
+            print(f"  Task {tid:3d}: {best[1]:40s} ({len(models[tid]):>8,} → {len(best[0]):>8,}) {score_s}")
+
+    print(f"\nDone in {time.time()-t0:.1f}s. S1:{counts[1]} S2:{counts[2]} S3:{counts[3]} S4:{counts[4]} Total:{len(results)}")
+
+    with zipfile.ZipFile(args.output_zip, 'w', zipfile.ZIP_DEFLATED) as zf:
+        for tid in range(1, 401):
+            zf.writestr(f'task{tid:03d}.onnx', results.get(tid, models[tid]))
+    print(f"Written to {args.output_zip}")
+
+
+if __name__ == '__main__':
+    main()