From 706569c57bfbdc7e9a0791cdb608236208012710 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Thu, 23 Mar 2023 09:24:59 +0100 Subject: [PATCH] Update --- beaver.py | 54 +++++++++++++++++++++++++++++------------------------- 1 file changed, 29 insertions(+), 25 deletions(-) diff --git a/beaver.py b/beaver.py index 8fe9a9b..69116ea 100755 --- a/beaver.py +++ b/beaver.py @@ -135,24 +135,33 @@ 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) - return order.sort(1).indices + order = order.sort(1).indices else: - return ( + order = ( torch.arange(x.size(1), device=x.device).unsqueeze(0).expand(x.size(0), -1) ) + return order -def shuffle(x, order, reorder=False): - if x.dim() == 3: - order = order.unsqueeze(-1).expand(-1, -1, x.size(-1)) - if reorder: - y = x.new(x.size()) - y.scatter_(1, order, x) - return y +def reorder(x, order, back=False): # x is NxTxD1x...xDk, order is NxT' + u = x.reshape(x.size()[:2] + (-1,)) + order = order.unsqueeze(-1).expand(-1, -1, u.size(-1)) + if back: + v = u.new(u.size()) + v.scatter_(1, order, u) else: - return x.gather(1, order) + v = u.gather(1, order) + v = v.reshape(v.size()[:2] + x.size()[2:]) + return v +def shuffle(x, fixed_len): + order = generation_order(x, fixed_len) + return reorder(x, order), order + + +###################################################################### + # ar_mask is a Boolean matrix of same shape as input, with 1s on the # tokens that should be generated @@ -186,8 +195,7 @@ def compute_perplexity(model, split="train"): for input in task.batches(split=split): input = input.to(device) - order = generation_order(input, task.height * task.width) - input = shuffle(input, order) + input, order = shuffle(input, task.height * task.width) output = model(mygpt.BracketedSequence(input), order=order).x loss = F.cross_entropy(output.transpose(1, 2), input) acc_loss += loss.item() * input.size(0) @@ -252,10 +260,9 @@ def oneshot(gpt, task): acc_train_loss, nb_train_samples = 0, 0 for mazes, policies in task.policy_batches(split="train"): - order = generation_order(mazes, task.height * task.width) - x = shuffle(mazes, order) + x, order = shuffle(mazes, task.height * task.width) x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x - output_gpt = shuffle(x, order, reorder=True) + output_gpt = reorder(x, order, back=True) output = model(output_gpt) loss = compute_loss(mazes, output, policies, task.height, task.width) @@ -268,10 +275,9 @@ def oneshot(gpt, task): acc_test_loss, nb_test_samples = 0, 0 for mazes, policies in task.policy_batches(split="test"): - order = generation_order(mazes, task.height * task.width) - x = shuffle(mazes, order) + x, order = shuffle(mazes, task.height * task.width) x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x - output_gpt = shuffle(x, order, reorder=True) + output_gpt = reorder(x, order, back=True) output = model(output_gpt) loss = compute_loss(mazes, output, policies, task.height, task.width) acc_test_loss += loss.item() * mazes.size(0) @@ -284,10 +290,9 @@ def oneshot(gpt, task): # ------------------- mazes = task.test_input[:32, : task.height * task.width] policies = task.test_policies[:32] - order = generation_order(mazes, task.height * task.width) - x = shuffle(mazes, order) + x, order = shuffle(mazes, task.height * task.width) x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x - output_gpt = shuffle(x, order, reorder=True) + output_gpt = reorder(x, order, back=True) output = model(output_gpt) if args.oneshot_output == "policy": targets = policies.permute(0, 2, 1) @@ -426,11 +431,11 @@ class TaskMaze(Task): ar_mask = result.new_zeros(result.size()) ar_mask[:, self.height * self.width :] = 1 result *= 1 - ar_mask - order = generation_order(result, self.height * self.width) + result, order = shuffle(result, self.height * self.width) masked_inplace_autoregression( model, self.batch_size, result, ar_mask, order=order ) - result = shuffle(result, order, reorder=True) + result = reorder(result, order, back=True) mazes, paths = self.seq2map(result) nb_correct += maze.path_correctness(mazes, paths).long().sum() nb_total += mazes.size(0) @@ -608,8 +613,7 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs): for input in task.batches(split="train"): input = input.to(device) - order = generation_order(input, task.height * task.width) - input = shuffle(input, order) + input, order = shuffle(input, task.height * task.width) output = model(mygpt.BracketedSequence(input), order=order).x loss = F.cross_entropy(output.transpose(1, 2), input) acc_train_loss += loss.item() * input.size(0) -- 2.39.5