From beb6caeb2ffc8c5185bddbc07c4486b3cb9c8495 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Thu, 23 Mar 2023 12:41:55 +0100 Subject: [PATCH] Update --- beaver.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/beaver.py b/beaver.py index bd17365..f5f092b 100755 --- a/beaver.py +++ b/beaver.py @@ -134,7 +134,7 @@ for n in vars(args): def generation_order(x, fixed_len): if args.random_regression_order: order = torch.rand(x.size(), device=x.device) - order[:, :fixed_len] = torch.linspace(-2, -1, fixed_len, device=order.device) + order[:, :fixed_len] = torch.linspace(-2, -1, fixed_len, device=x.device) order = order.sort(1).indices else: order = ( @@ -167,7 +167,9 @@ def shuffle(x, fixed_len): def masked_inplace_autoregression(model, batch_size, input, ar_mask, order=None): - for input, ar_mask in zip(input.split(batch_size), ar_mask.split(batch_size)): + for input, ar_mask, order in zip( + input.split(batch_size), ar_mask.split(batch_size), order.split(batch_size) + ): i = (ar_mask.sum(0) > 0).nonzero() if i.min() > 0: # Needed to initialize the model's cache @@ -186,7 +188,7 @@ def masked_inplace_autoregression(model, batch_size, input, ar_mask, order=None) ###################################################################### -def compute_perplexity(model, fixed_len, split="train"): +def compute_perplexity(model, task, fixed_len, split="train"): with torch.autograd.no_grad(): t = model.training model.eval() @@ -469,7 +471,11 @@ class TaskMaze(Task): ar_mask = result.new_zeros(result.size()) ar_mask[:, self.height * self.width :] = 1 result *= 1 - ar_mask - masked_inplace_autoregression(model, self.batch_size, result, ar_mask) + x, order = shuffle(result, self.height * self.width) + masked_inplace_autoregression( + model, self.batch_size, x, ar_mask, order=order + ) + result = reorder(x, order, back=True) mazes, paths = self.seq2map(input) _, predicted_paths = self.seq2map(result) @@ -588,10 +594,10 @@ log_string(f"learning_rate_schedule {learning_rate_schedule}") if nb_epochs_finished >= args.nb_epochs: n_epoch = nb_epochs_finished train_perplexity = compute_perplexity( - model, fixed_len=task.height * task.width, split="train" + model, task, fixed_len=task.height * task.width, split="train" ) test_perplexity = compute_perplexity( - model, fixed_len=task.height * task.width, split="test" + model, task, fixed_len=task.height * task.width, split="test" ) log_string( @@ -635,7 +641,7 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs): train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples)) test_perplexity = compute_perplexity( - model, fixed_len=task.height * task.width, split="test" + model, task, fixed_len=task.height * task.width, split="test" ) log_string( -- 2.39.5