X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=beaver.py;h=9f8bc4104e92189bcf76e265b32a23aa71c0ab2c;hb=763e18680a0ae2db64d27c9f2a2054a7403009c0;hp=6ed9dd23a384984374e8feff2ba489d103150606;hpb=92a2935401f6fe21efe19ac3c476521665a242a5;p=beaver.git

diff --git a/beaver.py b/beaver.py
index 6ed9dd2..9f8bc41 100755
--- a/beaver.py
+++ b/beaver.py
@@ -238,6 +238,11 @@ def oneshot(model, learning_rate_scheduler, task):
     model.eval()
     mazes = task.test_input[:32].clone()
     mazes[:, task.height * task.width :] = 0
+    policies = task.test_policies[:32]
+    targets = maze.stationary_densities(
+        mazes[:, : task.height * task.width].view(-1, task.height, task.width),
+        policies.view(-1, 4, task.height, task.width),
+    ).flatten(-2)
     output = eval_mygpt(model, mazes, prompt_len=task.height * task.width)
     output = F.softmax(output, dim=2)
     print(f"{output.size()=}")
@@ -245,13 +250,17 @@ def oneshot(model, learning_rate_scheduler, task):
         -1, task.height, task.width
     )
     mazes = mazes[:, : task.height * task.width].reshape(-1, task.height, task.width)
-    # targets = targets.reshape(-1, task.height, task.width)
+    targets = targets.reshape(-1, task.height, task.width)
+    paths = task.test_input[:32, task.height * task.width :].reshape(
+        -1, task.height, task.width
+    )
     filename = f"oneshot.png"
     maze.save_image(
         os.path.join(args.result_dir, filename),
         mazes=mazes,
+        # target_paths=paths,
         score_paths=proba_path,
-        # score_truth=targets,
+        score_truth=targets,
     )
     log_string(f"wrote {filename}")