Add --log_every arg to control epoch print frequency (default 100)

trm_solver/trm_nano/train_task.py

@@ -8,10 +8,11 @@ via manual builder (weight reuse), logs everything.
 Usage (Kaggle):
     python train_task.py --data_dir /kaggle/input/competitions/neurogolf-2026/ \
                          --output_dir /kaggle/working/trm_onnx/ \
-                         --skip_solved /kaggle/working/solved_by_analytical.json
+                         --skip_solved /kaggle/working/solved_by_analytical.json \
+                         --log_every 50
 
 Single task mode:
-    python train_task.py --task_file /path/to/task001.json --output_dir onnx_out/
+    python train_task.py --task_file /path/to/task001.json --output_dir onnx_out/ --log_every 25
 """
 
 import os
@@ -140,7 +141,7 @@ class ARCTaskDataset(Dataset):
 # ─── Training ─────────────────────────────────────────────────────────────────
 
 def train_on_task(task_data, config, epochs=500, lr=1e-3, batch_size=64,
-                  num_aug=1000, device="cuda"):
+                  num_aug=1000, device="cuda", log_every=100):
     model = TRMNano(config).to(device)
     dataset = ARCTaskDataset(task_data, num_augmentations=num_aug)
     loader = DataLoader(dataset, batch_size=batch_size, shuffle=True,
@@ -182,7 +183,7 @@ def train_on_task(task_data, config, epochs=500, lr=1e-3, batch_size=64,
             best_loss = avg_loss
             best_state = {n: p.clone() for n, p in ema_params.items()}
 
-        if (epoch + 1) % 100 == 0:
+        if (epoch + 1) % log_every == 0:
             logging.info(f"    epoch {epoch+1}/{epochs} loss={avg_loss:.4f} best={best_loss:.4f}")
 
     if best_state:
@@ -212,7 +213,6 @@ def verify_with_onnx(onnx_path, task_data):
             inp_grid = np.array(example["input"], dtype=np.int64)
             out_grid = np.array(example["output"], dtype=np.int64)
 
-            # Convert input to one-hot [1, 10, 30, 30]
             h, w = inp_grid.shape
             if h > 30 or w > 30:
                 fails += 1
@@ -222,11 +222,9 @@ def verify_with_onnx(onnx_path, task_data):
                 for c in range(w):
                     inp_onehot[0, inp_grid[r, c], r, c] = 1.0
 
-            # Run inference
             result = sess.run(None, {"input": inp_onehot})
             pred_onehot = (result[0] > 0.0).astype(np.float32)
 
-            # Convert expected output to one-hot
             oh, ow = out_grid.shape
             exp_onehot = np.zeros((1, 10, 30, 30), dtype=np.float32)
             for r in range(oh):
@@ -271,6 +269,8 @@ def main():
     parser.add_argument("--hidden_size", type=int, default=96)
     parser.add_argument("--H_cycles", type=int, default=3)
     parser.add_argument("--L_cycles", type=int, default=4)
+    parser.add_argument("--log_every", type=int, default=100,
+                        help="Print training loss every N epochs (default: 100)")
     # Limits
     parser.add_argument("--max_tasks", type=int, default=None)
     parser.add_argument("--start_task", type=int, default=1)
@@ -293,6 +293,7 @@ def main():
     logging.info(f"Params: {count_params(TRMNano(config)):,}")
     logging.info(f"Device: {args.device}")
     logging.info(f"Epochs: {args.epochs}, LR: {args.lr}, Batch: {args.batch_size}, Aug: {args.num_aug}")
+    logging.info(f"Log every: {args.log_every} epochs")
     logging.info(f"Output: {args.output_dir}")
 
     # ─── Single task mode ─────────────────────────────────────────────────────
@@ -305,7 +306,7 @@ def main():
         t0 = time.time()
         model, best_loss = train_on_task(task_data, config, epochs=args.epochs, lr=args.lr,
                                          batch_size=args.batch_size, num_aug=args.num_aug,
-                                         device=args.device)
+                                         device=args.device, log_every=args.log_every)
         elapsed = time.time() - t0
         logging.info(f"  Trained in {elapsed:.1f}s, best_loss={best_loss:.4f}")
 
@@ -322,7 +323,6 @@ def main():
         return
 
     # ─── Batch mode ───────────────────────────────────────────────────────────
-    # Load skip list
     skip_tasks = set()
     if args.skip_solved and os.path.exists(args.skip_solved):
         with open(args.skip_solved) as f:
@@ -333,26 +333,23 @@ def main():
                 skip_tasks = set(skip_data.keys())
         logging.info(f"Skipping {len(skip_tasks)} already-solved tasks")
 
-    # Find task files
     task_files = sorted(glob.glob(os.path.join(args.data_dir, "task*.json")))
     if not task_files:
         task_files = sorted(glob.glob(os.path.join(args.data_dir, "*.json")))
     logging.info(f"Found {len(task_files)} task files in {args.data_dir}")
 
-    # Filter
     tasks_to_run = []
     for f in task_files:
         task_id = os.path.basename(f).replace(".json", "")
-        # Extract task number if possible
         task_num = None
         try:
             task_num = int(task_id.replace("task", ""))
         except ValueError:
             pass
-
-        # Skip if in solved list (check both task_id and task_num)
         if task_id in skip_tasks or str(task_num) in skip_tasks:
             continue
+        if task_num is not None and task_num in skip_tasks:
+            continue
         if task_num is not None and task_num < args.start_task:
             continue
         tasks_to_run.append((task_id, task_num, f))
@@ -384,14 +381,12 @@ def main():
             model, best_loss = train_on_task(
                 task_data, config, epochs=args.epochs, lr=args.lr,
                 batch_size=args.batch_size, num_aug=args.num_aug,
-                device=args.device)
+                device=args.device, log_every=args.log_every)
             train_time = time.time() - t0
 
-            # Export ONNX
             onnx_path = os.path.join(args.output_dir, f"{task_id}.onnx")
             size_bytes = export_onnx_manual(model, config, onnx_path)
 
-            # Verify
             passes, fails = verify_with_onnx(onnx_path, task_data)
             solved = (fails == 0 and passes > 0)
             if solved:
@@ -411,7 +406,6 @@ def main():
                 "train_time_s": train_time,
             })
 
-            # Remove ONNX if not solved (save disk)
             if not solved and os.path.exists(onnx_path):
                 os.remove(onnx_path)
 
@@ -424,7 +418,6 @@ def main():
                 "error": str(e),
             })
 
-        # Save progress after each task
         progress_path = os.path.join(args.output_dir, "progress.json")
         with open(progress_path, "w") as f:
             json.dump({
@@ -445,7 +438,6 @@ def main():
     logging.info(f"  Output: {args.output_dir}")
     logging.info(f"  Log:    {log_path}")
 
-    # List solved ONNX files
     onnx_files = sorted(glob.glob(os.path.join(args.output_dir, "*.onnx")))
     total_kb = sum(os.path.getsize(f) for f in onnx_files) / 1024
     logging.info(f"  ONNX files: {len(onnx_files)} ({total_kb:.1f} KB total)")