From cd3fae5ff5e9b4c76b533dac8bc64cc331ecf51e Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Tue, 21 Mar 2023 06:33:47 +0100 Subject: [PATCH] Update --- beaver.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/beaver.py b/beaver.py index 5916215..f5bd924 100755 --- a/beaver.py +++ b/beaver.py @@ -79,6 +79,9 @@ 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") @@ -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) @@ -535,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") @@ -608,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) + +###################################################################### -- 2.39.5