Add sign_corrected_conv.py: refit conv_var tasks with direct output (>0 validation, 3 nodes instead of 15-20)

own-solver/sign_corrected_conv.py

@@ -0,0 +1,290 @@
+#!/usr/bin/env python3
+"""
+Sign-corrected direct Conv solver for conv_var tasks.
+
+Replaces the current architecture:
+  ReduceSum → Conv → ArgMax → Equal×10 → Cast×10 → Concat → Mul → output (15-20 nodes)
+
+With minimal architecture:
+  ReduceSum → Conv → Mul → output (3 nodes, 2 intermediates)
+
+The key insight: Kaggle validates via (output > 0.0).astype(float).
+If Conv weights are trained to produce POSITIVE values only in the correct
+channel and NEGATIVE/ZERO in all others, we skip ArgMax+OneHot entirely.
+
+Algorithm:
+1. lstsq fit with targets: +1 (correct channel), -1 (wrong channels)
+2. Check for sign violations (pixels where wrong channel is positive)
+3. Iteratively increase margin on violating pixels until all signs are correct
+4. Validate on full train+test+arc-gen
+
+Integration: Add this to conv.py as an alternative to solve_conv_variable.
+Call it FIRST (before the standard approach). If it succeeds, use it.
+If not, fall back to the standard ArgMax+OneHot approach.
+
+Score impact under new formula (25 - ln(memory + params)):
+  Old conv_var ks=3: ~10 intermediate tensors → cost ~170,000 → score ~13.0
+  New direct ks=3: 2 intermediate tensors → cost ~44,500 → score ~14.3
+  GAIN: +1.3 pts per ks=3 task (8 tasks = +10.4 pts)
+  
+  For larger kernels: gain is similar since we eliminate the same onehot overhead.
+
+Usage standalone:
+  python sign_corrected_conv.py --data_dir ./tasks --output_dir ./models
+
+Usage integrated in solver pipeline:
+  Import solve_conv_var_direct() and call before solve_conv_variable()
+"""
+
+import json
+import math
+import os
+import sys
+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]
+GH, GW = 30, 30
+
+
+def encode_grid(grid):
+    """Encode integer 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):
+            if 0 <= arr[r, c] < 10:
+                t[0, arr[r, c], r, c] = 1.0
+    return t
+
+
+def validate_model(model_bytes, examples, max_check=50):
+    """Validate model via (output > 0.0) == expected one-hot."""
+    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[:max_check]:
+        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
+
+
+def build_conv_var_direct(W, ks):
+    """Build minimal conv_var model: ReduceSum→Conv→Mul→output (3 nodes)."""
+    pad_k = ks // 2
+    w_init = numpy_helper.from_array(W.astype(np.float32), 'W')
+    ax_init = numpy_helper.from_array(np.array([1], dtype=np.int64), 'ax')
+    nodes = [
+        helper.make_node('ReduceSum', ['input', 'ax'], ['mask'], keepdims=1),
+        helper.make_node('Conv', ['input', 'W'], ['co'],
+                        kernel_shape=[ks, ks], pads=[pad_k, pad_k, pad_k, pad_k]),
+        helper.make_node('Mul', ['co', 'mask'], ['output']),
+    ]
+    x = helper.make_tensor_value_info('input', TensorProto.FLOAT, GRID_SHAPE)
+    y = helper.make_tensor_value_info('output', TensorProto.FLOAT, GRID_SHAPE)
+    g = helper.make_graph(nodes, 'g', [x], [y], initializer=[w_init, ax_init])
+    return helper.make_model(g, ir_version=8, opset_imports=[helper.make_opsetid('', 17)])
+
+
+def fit_sign_corrected_conv(examples, ks, max_iter=30, use_full_30=True):
+    """Fit Conv weights with iterative sign correction.
+    
+    Returns W (10, 10, ks, ks) float32 or None if fitting fails.
+    
+    The fitted weights satisfy: for every training pixel,
+    Conv output is POSITIVE only in the correct color channel
+    and NEGATIVE/ZERO in all other channels.
+    """
+    pad_k = ks // 2
+    
+    # Build patch matrix from ground truth examples
+    patches = []
+    target_classes = []
+    
+    for ex in examples:
+        inp_grid = np.array(ex['input'])
+        out_grid = np.array(ex['output'])
+        ih, iw = inp_grid.shape
+        oh, ow = out_grid.shape
+        
+        # One-hot encode input on full 30×30 grid
+        inp_oh = np.zeros((10, 30, 30), dtype=np.float64)
+        for c in range(10):
+            inp_oh[c, :ih, :iw] = (inp_grid == c)
+        inp_padded = np.pad(inp_oh, ((0, 0), (pad_k, pad_k), (pad_k, pad_k)))
+        
+        for r in range(oh):
+            for c in range(ow):
+                patch = inp_padded[:, r:r + ks, c:c + ks].flatten()
+                patches.append(patch)
+                target_classes.append(int(out_grid[r, c]))
+    
+    if not patches:
+        return None
+    
+    P = np.array(patches, dtype=np.float64)
+    tc = np.array(target_classes, dtype=np.int64)
+    n = len(P)
+    feat = P.shape[1]
+    
+    # Skip if problem is too large
+    if feat > 5000 and n > 5000:
+        return None
+    
+    # Initial targets: +1 for correct class, -1 for others
+    T = np.full((n, 10), -1.0, dtype=np.float64)
+    for i in range(n):
+        if 0 <= tc[i] < 10:
+            T[i, tc[i]] = 1.0
+    
+    for it in range(max_iter):
+        try:
+            W_flat = np.linalg.lstsq(P, T, rcond=None)[0]
+        except (np.linalg.LinAlgError, ValueError):
+            return None
+        
+        pred = P @ W_flat
+        
+        # Count and fix sign violations
+        violations = 0
+        for i in range(n):
+            c = tc[i]
+            if c < 0 or c >= 10:
+                continue
+            # Correct class must be > 0
+            if pred[i, c] <= 0:
+                violations += 1
+                T[i, c] = 2.0 + it * 2.0  # Increase positive margin
+            # Wrong classes must be <= 0
+            for j in range(10):
+                if j != c and pred[i, j] > 0:
+                    violations += 1
+                    T[i, j] = -(2.0 + it * 2.0)  # Increase negative margin
+        
+        if violations == 0:
+            W = W_flat.T.reshape(10, 10, ks, ks).astype(np.float32)
+            return W
+    
+    return None  # Failed to eliminate all violations
+
+
+def solve_conv_var_direct(td, ks_list=None):
+    """Attempt to solve a task using sign-corrected direct Conv.
+    
+    Args:
+        td: task data dict with 'train', 'test', 'arc-gen' keys
+        ks_list: kernel sizes to try (default: [3, 5, 7])
+    
+    Returns:
+        (model_bytes, ks) on success, None on failure.
+    """
+    if ks_list is None:
+        ks_list = [3, 5, 7, 9, 11, 13, 15]
+    
+    train_examples = td.get('train', []) + td.get('test', [])
+    arcgen = td.get('arc-gen', [])[:30]
+    all_examples = train_examples + arcgen
+    
+    if len(train_examples) < 2:
+        return None
+    
+    # Check all examples have same input/output shape (variable size OK)
+    for ex in train_examples:
+        inp = np.array(ex['input'])
+        out = np.array(ex['output'])
+        if inp.shape != out.shape:
+            return None  # Different shapes not supported by this approach
+    
+    for ks in ks_list:
+        # Skip oversized problems
+        n_pixels = sum(np.array(ex['output']).size for ex in train_examples)
+        feat = 10 * ks * ks
+        if feat > 3000 and n_pixels > 3000:
+            continue
+        
+        W = fit_sign_corrected_conv(train_examples, ks, max_iter=25)
+        if W is None:
+            continue
+        
+        # Build and validate
+        model = build_conv_var_direct(W, ks)
+        model_bytes = model.SerializeToString()
+        
+        if validate_model(model_bytes, all_examples):
+            return model_bytes, ks
+    
+    return None
+
+
+def main():
+    """Standalone: process all tasks in data_dir, save models to output_dir."""
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data_dir', required=True, help='Directory with taskNNN.json')
+    parser.add_argument('--output_dir', required=True, help='Output directory for .onnx models')
+    parser.add_argument('--tasks', type=str, default='',
+                       help='Comma-separated task IDs (default: all conv_var tasks)')
+    args = parser.parse_args()
+    
+    os.makedirs(args.output_dir, exist_ok=True)
+    
+    # Default: known conv_var tasks from V90
+    if args.tasks:
+        task_ids = [int(t) for t in args.tasks.split(',')]
+    else:
+        task_ids = [120, 127, 171, 220, 230, 258, 261, 352,  # ks=3
+                    389,  # ks=5
+                    32,   # ks=9
+                    82, 122, 288, 305,  # ks=11
+                    359,  # ks=13
+                    28, 237, 348]  # ks=15
+    
+    print(f"Processing {len(task_ids)} tasks")
+    successes = 0
+    
+    for tid in task_ids:
+        task_path = os.path.join(args.data_dir, f'task{tid:03d}.json')
+        if not os.path.exists(task_path):
+            print(f"  Task {tid:3d}: SKIPPED (no data file)")
+            continue
+        
+        with open(task_path) as f:
+            td = json.load(f)
+        
+        result = solve_conv_var_direct(td)
+        if result:
+            model_bytes, ks = result
+            out_path = os.path.join(args.output_dir, f'task{tid:03d}.onnx')
+            with open(out_path, 'wb') as f:
+                f.write(model_bytes)
+            
+            # Estimate score
+            W_elements = 10 * 10 * ks * ks
+            cost = W_elements * 5 + 39609  # weight_mem + intermediates + params
+            score = max(1.0, 25.0 - math.log(cost))
+            
+            print(f"  Task {tid:3d}: ✅ ks={ks} ({len(model_bytes):,} bytes) score≈{score:.2f}")
+            successes += 1
+        else:
+            print(f"  Task {tid:3d}: ❌ sign correction failed")
+    
+    print(f"\nDone: {successes}/{len(task_ids)} tasks converted to direct Conv")
+
+
+if __name__ == '__main__':
+    main()