v4.2: Add PyTorch learned conv solver (single+two-layer, multi-seed, ternary snap). Needs GPU to be practical - use on Kaggle with --conv_budget 60

neurogolf_solver.py

@@ -1261,6 +1261,220 @@ def solve_conv_var_diff(td, path, time_budget=30.0):
             if validate(path, td): return 'conv_var_diff', model
     return None
 
+# ============================================================
+# PYTORCH LEARNED CONV (gradient descent, multi-seed, ternary snap)
+# ============================================================
+
+def _ternary_snap(w, eps=0.2):
+    """Snap weights to {-1, 0, 1} — smaller model, often still correct."""
+    return np.where(w > eps, 1.0, np.where(w < -eps, -1.0, 0.0)).astype(np.float32)
+
+def _build_conv_onnx_from_weights(W, ks, use_full_30=False, IH=None, IW=None):
+    """Build ONNX conv model from numpy weight array W [10,10,ks,ks].
+    For fixed-shape: Slice→Conv→ArgMax→Equal+Cast→Pad
+    For variable/full30: Conv→ArgMax→Equal+Cast→Mul(mask)"""
+    pad = ks // 2
+    if use_full_30:
+        # Variable shape: full 30x30 conv with mask
+        inits = [numpy_helper.from_array(W, 'W')]
+        nodes = [
+            helper.make_node('ReduceSum', ['input'], ['mask'], axes=[1], keepdims=1),
+            helper.make_node('Conv', ['input', 'W'], ['co'], kernel_shape=[ks,ks], pads=[pad]*4),
+            helper.make_node('ArgMax', ['co'], ['am'], axis=1, keepdims=1),
+        ]
+        add_onehot_block(nodes, inits, 'am', 'oh_out')
+        nodes.append(helper.make_node('Mul', ['oh_out', 'mask'], ['output']))
+        return mk(nodes, inits)
+    else:
+        # Fixed shape: slice, conv, pad
+        pad_h, pad_w = GH - IH, GW - IW
+        inits = [
+            numpy_helper.from_array(np.array([0,0,0,0], dtype=np.int64), 'sl_st'),
+            numpy_helper.from_array(np.array([1,10,IH,IW], dtype=np.int64), 'sl_en'),
+            numpy_helper.from_array(W, 'W'),
+        ]
+        nodes = [
+            helper.make_node('Slice', ['input','sl_st','sl_en'], ['grid']),
+            helper.make_node('Conv', ['grid', 'W'], ['co'], kernel_shape=[ks,ks], pads=[pad]*4),
+            helper.make_node('ArgMax', ['co'], ['am'], axis=1, keepdims=1),
+        ]
+        add_onehot_block(nodes, inits, 'am', 'oh_out')
+        nodes.append(
+            helper.make_node('Pad', ['oh_out'], ['output'],
+                pads=[0,0,0,0,0,0,pad_h,pad_w], value=0.0)
+        )
+        return mk(nodes, inits)
+
+def _build_two_layer_conv_onnx(W1, W2, ks1, ks2, use_full_30=False, IH=None, IW=None):
+    """Build ONNX two-layer conv: Conv→ReLU→Conv→ArgMax→Equal+Cast→Pad/Mul(mask)."""
+    pad1, pad2 = ks1 // 2, ks2 // 2
+    if use_full_30:
+        inits = [
+            numpy_helper.from_array(W1, 'W1'),
+            numpy_helper.from_array(W2, 'W2'),
+        ]
+        nodes = [
+            helper.make_node('ReduceSum', ['input'], ['mask'], axes=[1], keepdims=1),
+            helper.make_node('Conv', ['input', 'W1'], ['h1'], kernel_shape=[ks1,ks1], pads=[pad1]*4),
+            helper.make_node('Relu', ['h1'], ['h1r']),
+            helper.make_node('Conv', ['h1r', 'W2'], ['co'], kernel_shape=[ks2,ks2], pads=[pad2]*4),
+            helper.make_node('ArgMax', ['co'], ['am'], axis=1, keepdims=1),
+        ]
+        add_onehot_block(nodes, inits, 'am', 'oh_out')
+        nodes.append(helper.make_node('Mul', ['oh_out', 'mask'], ['output']))
+        return mk(nodes, inits)
+    else:
+        pad_h, pad_w = GH - IH, GW - IW
+        inits = [
+            numpy_helper.from_array(np.array([0,0,0,0], dtype=np.int64), 'sl_st'),
+            numpy_helper.from_array(np.array([1,10,IH,IW], dtype=np.int64), 'sl_en'),
+            numpy_helper.from_array(W1, 'W1'),
+            numpy_helper.from_array(W2, 'W2'),
+        ]
+        nodes = [
+            helper.make_node('Slice', ['input','sl_st','sl_en'], ['grid']),
+            helper.make_node('Conv', ['grid', 'W1'], ['h1'], kernel_shape=[ks1,ks1], pads=[pad1]*4),
+            helper.make_node('Relu', ['h1'], ['h1r']),
+            helper.make_node('Conv', ['h1r', 'W2'], ['co'], kernel_shape=[ks2,ks2], pads=[pad2]*4),
+            helper.make_node('ArgMax', ['co'], ['am'], axis=1, keepdims=1),
+        ]
+        add_onehot_block(nodes, inits, 'am', 'oh_out')
+        nodes.append(
+            helper.make_node('Pad', ['oh_out'], ['output'],
+                pads=[0,0,0,0,0,0,pad_h,pad_w], value=0.0)
+        )
+        return mk(nodes, inits)
+
+def solve_pytorch_conv(td, path, time_budget=30.0):
+    """PyTorch gradient descent conv solver. Tries single-layer then two-layer.
+    Multi-seed training with ternary weight snapping for smaller models.
+    Validates against arc-gen before accepting."""
+    try:
+        import torch
+        import torch.nn as nn
+        import copy as _copy
+    except ImportError:
+        return None
+    
+    exs = get_exs(td)
+    same_shape = all(inp.shape == out.shape for inp, out in exs)
+    if not same_shape:
+        return None  # Only handle same-shape for now
+    
+    shapes = set(inp.shape for inp, _ in exs)
+    fixed_in = len(shapes) == 1
+    
+    # Prepare tensors
+    all_pairs = td['train'] + td['test']
+    inp_list = [to_onehot(p['input'])[0] for p in all_pairs]
+    out_list = [to_onehot(p['output'])[0] for p in all_pairs]
+    inp_t = torch.tensor(np.stack(inp_list), dtype=torch.float32)
+    out_t = torch.tensor(np.stack(out_list), dtype=torch.float32)
+    
+    if fixed_in:
+        IH, IW = list(shapes)[0]
+        # Train on cropped region
+        inp_t = inp_t[:, :, :IH, :IW]
+        out_t = out_t[:, :, :IH, :IW]
+    
+    t_start = time.time()
+    best_result = None
+    
+    # Phase 1: Single-layer conv (multiple kernel sizes and seeds)
+    for ks in [1, 3, 5, 7]:
+        if time.time() - t_start > time_budget * 0.6:
+            break
+        pad = ks // 2
+        for seed in [0, 7, 42]:
+            if time.time() - t_start > time_budget * 0.6:
+                break
+            torch.manual_seed(seed)
+            conv = nn.Conv2d(CH, CH, kernel_size=ks, padding=pad, bias=False)
+            if seed == 0:
+                nn.init.zeros_(conv.weight)
+            opt = torch.optim.Adam(conv.parameters(), lr=0.03)
+            best_loss, best_state = float('inf'), None
+            for step in range(3000):
+                opt.zero_grad()
+                pred = conv(inp_t)
+                loss = nn.functional.mse_loss(pred, out_t)
+                loss.backward()
+                opt.step()
+                if loss.item() < best_loss:
+                    best_loss = loss.item()
+                    best_state = _copy.deepcopy(conv.state_dict())
+                if best_loss < 1e-8:
+                    break
+            if best_state is None:
+                continue
+            conv.load_state_dict(best_state)
+            w = conv.weight.detach().numpy()
+            
+            # Try continuous weights, then ternary-snapped
+            for w_cand in [w, _ternary_snap(w)]:
+                use_full = not fixed_in
+                model = _build_conv_onnx_from_weights(
+                    w_cand, ks, use_full_30=use_full,
+                    IH=IH if fixed_in else None,
+                    IW=IW if fixed_in else None
+                )
+                onnx.save(model, path)
+                if validate(path, td):
+                    sz = os.path.getsize(path)
+                    if best_result is None or sz < best_result[2]:
+                        best_result = ('pt_conv', model, sz)
+    
+    # Phase 2: Two-layer conv (Conv→ReLU→Conv)
+    for ks1, ks2, hidden in [(3, 1, CH), (5, 1, CH), (3, 3, CH)]:
+        if time.time() - t_start > time_budget:
+            break
+        for seed in [0, 7]:
+            if time.time() - t_start > time_budget:
+                break
+            torch.manual_seed(seed)
+            net = nn.Sequential(
+                nn.Conv2d(CH, hidden, kernel_size=ks1, padding=ks1//2, bias=False),
+                nn.ReLU(),
+                nn.Conv2d(hidden, CH, kernel_size=ks2, padding=ks2//2, bias=False),
+            )
+            opt = torch.optim.Adam(net.parameters(), lr=0.01)
+            best_loss, best_state = float('inf'), None
+            for step in range(2500):
+                opt.zero_grad()
+                pred = net(inp_t)
+                loss = nn.functional.mse_loss(pred, out_t)
+                loss.backward()
+                opt.step()
+                if loss.item() < best_loss:
+                    best_loss = loss.item()
+                    best_state = _copy.deepcopy(net.state_dict())
+                if best_loss < 1e-8:
+                    break
+            if best_state is None:
+                continue
+            net.load_state_dict(best_state)
+            w1 = net[0].weight.detach().numpy()
+            w2 = net[2].weight.detach().numpy()
+            
+            for w1c, w2c in [(w1, w2), (_ternary_snap(w1), _ternary_snap(w2))]:
+                use_full = not fixed_in
+                model = _build_two_layer_conv_onnx(
+                    w1c, w2c, ks1, ks2, use_full_30=use_full,
+                    IH=IH if fixed_in else None,
+                    IW=IW if fixed_in else None
+                )
+                onnx.save(model, path)
+                if validate(path, td):
+                    sz = os.path.getsize(path)
+                    if best_result is None or sz < best_result[2]:
+                        best_result = ('pt_conv2', model, sz)
+    
+    if best_result is not None:
+        sname, model, _ = best_result
+        onnx.save(model, path)
+        return sname, model
+    return None
+
 # ============================================================
 # MAIN
 # ============================================================
@@ -1316,6 +1530,12 @@ def solve_task(tn, td, outdir, conv_budget=30.0):
         if result is not None:
             sname, model = result
             return True, sname, os.path.getsize(path), time.time() - t_start, path
+        # 3. PyTorch learned conv as fallback for same-shape tasks
+        remaining = max(1, conv_time - (time.time() - t_start))
+        result = solve_pytorch_conv(td, path, time_budget=remaining)
+        if result is not None:
+            sname, model = result
+            return True, sname, os.path.getsize(path), time.time() - t_start, path
     else:
         sp = fixed_shapes(td)
         if sp is not None: