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

diff --git a/beaver.py b/beaver.py
index 6e1eaf4..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,6 +250,7 @@ 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)
     paths = task.test_input[:32, task.height * task.width :].reshape(
         -1, task.height, task.width
     )
@@ -252,9 +258,9 @@ def oneshot(model, learning_rate_scheduler, task):
     maze.save_image(
         os.path.join(args.result_dir, filename),
         mazes=mazes,
-        target_paths=paths,
+        # target_paths=paths,
         score_paths=proba_path,
-        # score_truth=targets,
+        score_truth=targets,
     )
     log_string(f"wrote {filename}")