X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=beaver.py;h=bdc12aa405624a866de0f9f2eb320f5aedf53210;hb=041605103d6529e5c03fc8ffa98a9a81a78842fb;hp=a3a56151f2ddb18994580979c5a4c568611108f5;hpb=15d0ad06d29c3e8042bdf64b78a01572e9b7ddc5;p=beaver.git

diff --git a/beaver.py b/beaver.py
index a3a5615..bdc12aa 100755
--- a/beaver.py
+++ b/beaver.py
@@ -79,11 +79,14 @@ parser.add_argument("--maze_width", type=int, default=21)
 
 parser.add_argument("--maze_nb_walls", type=int, default=15)
 
+##############################
+# one-shot prediction
+
 parser.add_argument("--oneshot", action="store_true", default=False)
 
 parser.add_argument("--oneshot_input", type=str, default="head")
 
-parser.add_argument("--oneshot_output", type=str, default="policy")
+parser.add_argument("--oneshot_output", type=str, default="trace")
 
 ######################################################################
 
@@ -224,11 +227,6 @@ def oneshot(gpt, task):
 
         acc_train_loss, nb_train_samples = 0, 0
         for mazes, policies in task.policy_batches(split="train"):
-            ####
-            # print(f'{mazes.size()=} {policies.size()=}')
-            # s = maze.stationary_densities(
-            # exit(0)
-            ####
             output_gpt = gpt(mygpt.BracketedSequence(mazes), mode=args.oneshot_input).x
             output = model(output_gpt)
 
@@ -268,7 +266,6 @@ def oneshot(gpt, task):
                 policies.view(-1, 4, task.height, task.width),
             ).flatten(-2)
             scores = output
-            print(scores)
         else:
             raise ValueError(f"{args.oneshot_output=}")
 
@@ -536,12 +533,6 @@ log_string(f"learning_rate_schedule {learning_rate_schedule}")
 
 ##############################
 
-if args.oneshot:
-    oneshot(model, task)
-    exit(0)
-
-##############################
-
 if nb_epochs_finished >= args.nb_epochs:
     n_epoch = nb_epochs_finished
     train_perplexity = compute_perplexity(model, split="train")
@@ -609,3 +600,8 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs):
     log_string(f"saved checkpoint {checkpoint_name}")
 
 ######################################################################
+
+if args.oneshot:
+    oneshot(model, task)
+
+######################################################################