Update neurogolf_solver.py

neurogolf_solver.py

@@ -976,7 +976,9 @@ def solve_task(tn, td, outdir, conv_budget=30.0):
 
 def run_tasks(task_nums, tasks, output_dir, conv_budget, use_wandb):
     results = {}
-
+    cost = 0
+    score = 0
+    total_score = 0
     for tn in task_nums:
         if tn not in tasks:
             continue
@@ -991,13 +993,15 @@ def run_tasks(task_nums, tasks, output_dir, conv_budget, use_wandb):
                     macs, memory, params = 0, 0, 0
             except:
                 macs, memory, params = 0, 0, 0
-            score = macs + memory + params
+            cost = macs + memory + params
+            score = max(1, 25 - math.log(max(1, cost)))
+            total_score += score
 
             results[tn] = (sname, t_task, sz)
-            print(f"Task {tn:3d}: {sname:25s} {score:>12} {t_task:7.3f}s  ({sz:>8,} bytes)")
+            print(f"Task {tn:3d}: {sname:25s} {score:7.3f} {cost:>12} {t_task:7.3f}s  ({sz:>8,} bytes)")
         else:
             print(f"Task {tn:3d}: UNSOLVED  {t_task:7.3f}s")
-            macs, memory, params, score = 0, 0, 0, 0
+            macs, memory, params, cost = 0, 0, 0, 0
 
         if use_wandb and wandb is not None:
             wandb.log({
@@ -1008,10 +1012,11 @@ def run_tasks(task_nums, tasks, output_dir, conv_budget, use_wandb):
                 "macs": macs,
                 "memory": memory,
                 "params": params,
+                "cost": cost,
                 "score": score,
             })
 
-    return results
+    return results, total_score
 
 
 def main():
@@ -1052,9 +1057,9 @@ def main():
             name="solver_run",
             config=config,
         ):
-            results = run_tasks(task_nums, tasks, args.output_dir, args.conv_budget, use_wandb=True)
+            results, total_score = run_tasks(task_nums, tasks, args.output_dir, args.conv_budget, use_wandb=True)
     else:
-        results = run_tasks(task_nums, tasks, args.output_dir, args.conv_budget, use_wandb=False)
+        results, total_score = run_tasks(task_nums, tasks, args.output_dir, args.conv_budget, use_wandb=False)
 
     elapsed = time.time() - t0
     print(f"\n{'='*70}")
@@ -1065,7 +1070,7 @@ def main():
     n_files = len([f for f in os.listdir(args.output_dir) if f.endswith('.onnx')])
     total_size = sum(os.path.getsize(os.path.join(args.output_dir, f))
                      for f in os.listdir(args.output_dir) if f.endswith('.onnx'))
-    print(f"\n{n_files} ONNX files, total {total_size/1024:.1f} KB")
+    print(f"\n{n_files} ONNX files, Total local estimated score: {total_score:.1f} total {total_size/1024:.1f} KB")
 
 if __name__ == '__main__':
     main()