Add Stage 1 trivial task optimizer (zero-intermediate architectures for max score)

own-solver/stage1_trivial_optimizer.py

@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+"""
+Stage 1: Rebuild trivial tasks for maximum score under new formula.
+
+New formula: score = max(1.0, 25.0 - ln(memory + params))
+Where memory = sum of ALL intermediate tensor bytes (excluding 'input'/'output')
+      params = sum of all initializer element counts + Constant node values
+
+KEY INSIGHT: If a node's output IS the graph output ('output'), it's NOT counted.
+So: input → [single op] → output = ZERO intermediate memory cost.
+Only weights/constants contribute.
+
+This script rewrites existing ONNX models to use minimal architectures:
+- Transpose: input→Transpose→output (cost=0, score=25)  
+- Color permutation: input→Gather(axis=1)→output (cost=40+10=50, score=21)
+- Color mapping: input→Conv1x1→output (cost=500, score=18.8)
+- Identity: input→Identity→output (cost=0, score=25)
+- Flips: input→Slice(reverse)→output (cost=32+4=36, score=21.4)
+
+Usage:
+  python stage1_trivial_optimizer.py --input_zip submission.zip --data_dir ./tasks --output_zip submission_stage1.zip
+"""
+
+import json
+import math
+import os
+import sys
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import onnx
+import onnxruntime as ort
+from onnx import helper, TensorProto, numpy_helper
+
+# ─── Config ───
+GRID_SHAPE = [1, 10, 30, 30]
+DT = TensorProto.FLOAT
+IR = 8
+
+
+def make_model(nodes, inits=None, opset=17):
+    """Create minimal ONNX model."""
+    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):
+    """Encode grid to one-hot tensor."""
+    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):
+    """Validate model produces correct output on all examples."""
+    try:
+        opts = ort.SessionOptions()
+        opts.log_severity_level = 3
+        sess = ort.InferenceSession(model_bytes, sess_options=opts, providers=['CPUExecutionProvider'])
+    except Exception:
+        return False
+    for ex in examples:
+        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(float), (expected > 0.0).astype(float)):
+                return False
+        except Exception:
+            return False
+    return True
+
+
+# ─── Stage 1 Optimizers ───
+
+def optimize_transpose(model_bytes):
+    """Rebuild as: input → Transpose → output (cost=0, score=25)."""
+    nodes = [helper.make_node('Transpose', ['input'], ['output'], perm=[0, 1, 3, 2])]
+    model = make_model(nodes)
+    return model.SerializeToString()
+
+
+def optimize_identity(model_bytes):
+    """Rebuild as: input → Identity → output (cost=0, score=25)."""
+    nodes = [helper.make_node('Identity', ['input'], ['output'])]
+    model = make_model(nodes)
+    return model.SerializeToString()
+
+
+def optimize_color_permutation(model_bytes, perm):
+    """Rebuild as: input → Gather(axis=1) → output.
+    perm[i] = source channel for output channel i.
+    Cost: 10 int32 elements = 40 bytes memory + 10 params = 50. Score ≈ 21.1
+    """
+    gi = np.array(perm, dtype=np.int32)
+    inits = [numpy_helper.from_array(gi, 'gi')]
+    nodes = [helper.make_node('Gather', ['input', 'gi'], ['output'], axis=1)]
+    model = make_model(nodes, inits)
+    return model.SerializeToString()
+
+
+def optimize_color_map_conv1x1(model_bytes, color_map):
+    """Rebuild as: input → Conv(1x1) → output.
+    Cost: 100 float32 elements = 400 bytes + 100 params = 500. Score ≈ 18.8
+    """
+    W = np.zeros((10, 10, 1, 1), dtype=np.float32)
+    for ic in range(10):
+        oc = color_map.get(ic, ic)
+        if 0 <= oc < 10:
+            W[oc, ic, 0, 0] = 1.0
+    inits = [numpy_helper.from_array(W, 'W')]
+    nodes = [helper.make_node('Conv', ['input', 'W'], ['output'], kernel_shape=[1, 1])]
+    model = make_model(nodes, inits)
+    return model.SerializeToString()
+
+
+def optimize_flip(axis):
+    """Rebuild as: input → Slice(reverse) → output.
+    Cost: 4 int64 scalars = 32 bytes + 4 params = 36. Score ≈ 21.4
+    """
+    starts = np.array([29], dtype=np.int64)
+    ends = np.array([np.iinfo(np.int64).min], dtype=np.int64)
+    axes = np.array([axis], dtype=np.int64)
+    steps = np.array([-1], dtype=np.int64)
+    inits = [
+        numpy_helper.from_array(starts, 'st'),
+        numpy_helper.from_array(ends, 'en'),
+        numpy_helper.from_array(axes, 'ax'),
+        numpy_helper.from_array(steps, 'sp'),
+    ]
+    nodes = [helper.make_node('Slice', ['input', 'st', 'en', 'ax', 'sp'], ['output'])]
+    model = make_model(nodes, inits)
+    return model.SerializeToString()
+
+
+def detect_and_optimize(task_id, model_bytes, examples):
+    """Detect if task can be optimized with a simpler architecture.
+    Returns (optimized_bytes, name, estimated_score) or None.
+    """
+    # Try identity
+    model_id = optimize_identity(model_bytes)
+    if validate_model(model_id, examples):
+        return model_id, "identity_direct", 25.0
+    
+    # Try transpose
+    model_t = optimize_transpose(model_bytes)
+    if validate_model(model_t, examples):
+        return model_t, "transpose_direct", 25.0
+    
+    # Try flips
+    for axis, name in [(3, 'flip_lr'), (2, 'flip_ud')]:
+        opt = optimize_flip(axis)
+        if validate_model(opt, examples):
+            return opt, f"{name}_direct", 21.4
+    
+    # Try color map detection
+    cm = {}
+    is_color_map = True
+    for ex in examples:
+        inp, out = np.array(ex['input']), np.array(ex['output'])
+        if inp.shape != out.shape:
+            is_color_map = 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_color_map = False
+                break
+            cm[iv] = ov
+        if not is_color_map:
+            break
+    
+    if is_color_map 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
+            opt = optimize_color_permutation(model_bytes, gather_ch)
+            if validate_model(opt, examples):
+                return opt, "color_perm_direct", 21.1
+        
+        opt = optimize_color_map_conv1x1(model_bytes, cm)
+        if validate_model(opt, examples):
+            return opt, "color_map_conv1x1_direct", 18.8
+    
+    return None
+
+
+def main():
+    """Process submission zip and optimize trivial tasks."""
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input_zip', required=True, help='Input submission.zip')
+    parser.add_argument('--data_dir', required=True, help='Directory with taskNNN.json files')
+    parser.add_argument('--output_zip', required=True, help='Output optimized submission.zip')
+    args = parser.parse_args()
+    
+    # Load all models from zip
+    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 from {args.input_zip}")
+    
+    # Process each task
+    optimized = {}
+    total_score_gain = 0.0
+    
+    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', [])
+        arcgen = task_data.get('arc-gen', [])[:30]
+        all_examples = examples + arcgen
+        
+        if not all_examples:
+            continue
+        
+        result = detect_and_optimize(tid, models[tid], all_examples)
+        if result:
+            opt_bytes, opt_name, est_score = result
+            orig_size = len(models[tid])
+            opt_size = len(opt_bytes)
+            optimized[tid] = opt_bytes
+            print(f"  Task {tid:3d}: {opt_name:30s} ({orig_size:>6,} → {opt_size:>6,} bytes) est_score={est_score:.1f}")
+    
+    print(f"\nOptimized {len(optimized)} tasks (Stage 1: trivial rebuilds)")
+    
+    # Write output zip
+    with zipfile.ZipFile(args.output_zip, 'w', zipfile.ZIP_DEFLATED) as zf:
+        for tid in range(1, 401):
+            fname = f'task{tid:03d}.onnx'
+            if tid in optimized:
+                zf.writestr(fname, optimized[tid])
+            elif tid in models:
+                zf.writestr(fname, models[tid])
+    
+    print(f"Written to {args.output_zip}")
+
+
+if __name__ == '__main__':
+    main()