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")
######################################################################
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)
- ####
- masks = mazes == maze.v_empty
output_gpt = gpt(mygpt.BracketedSequence(mazes), mode=args.oneshot_input).x
output = model(output_gpt)
mazes.view(-1, task.height, task.width),
policies.view(-1, 4, task.height, task.width),
).flatten(-2)
- scores = output.flatten(-2)
+ scores = output
else:
raise ValueError(f"{args.oneshot_output=}")
##############################
-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")
log_string(f"saved checkpoint {checkpoint_name}")
######################################################################
+
+if args.oneshot:
+ oneshot(model, task)
+
+######################################################################