From 88bcf05864ddf89d071ee4be17af57b3b3ce7c2a Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Thu, 23 Mar 2023 19:11:47 +0100 Subject: [PATCH] Update --- beaver.py | 41 ++++++++++++++++++++++------------------- 1 file changed, 22 insertions(+), 19 deletions(-) diff --git a/beaver.py b/beaver.py index 86008f6..6a6343d 100755 --- a/beaver.py +++ b/beaver.py @@ -131,10 +131,10 @@ for n in vars(args): ###################################################################### -def generation_order(x, fixed_len): +def generation_order(x, fixed_len=0): if args.random_regression_order: order = torch.rand(x.size(), device=x.device) - order[:, :fixed_len] = torch.linspace(-2, -1, fixed_len, device=x.device) + order[:, :fixed_len] = torch.arange(-fixed_len, 0, device=x.device) order = order.sort(1).indices else: order = ( @@ -147,8 +147,7 @@ def reorder(x, order, reverse=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 reverse: - v = u.new(u.size()) - v.scatter_(1, order, u) + v = u.new(u.size()).scatter_(1, order, u) else: v = u.gather(1, order) v = v.reshape(v.size()[:2] + x.size()[2:]) @@ -160,6 +159,12 @@ def shuffle(x, fixed_len): return reorder(x, order), order +def eval_mygpt(model, input, mode="standard", fixed_len=0): + x, order = shuffle(input, fixed_len) + x = model(mygpt.BracketedSequence(x), mode=mode, order=order).x + return reorder(x, order, reverse=True) + + ###################################################################### # ar_mask is a Boolean matrix of same shape as input, with 1s on the @@ -197,9 +202,7 @@ def compute_perplexity(model, task, fixed_len, split="train"): for input in task.batches(split=split): input = input.to(device) - x, order = shuffle(input, fixed_len) - x = model(mygpt.BracketedSequence(x), order=order).x - output = reorder(x, order, reverse=True) + output = eval_mygpt(model, input, fixed_len=fixed_len) loss = F.cross_entropy(output.transpose(1, 2), input) acc_loss += loss.item() * input.size(0) nb_samples += input.size(0) @@ -263,9 +266,9 @@ def oneshot(gpt, task): acc_train_loss, nb_train_samples = 0, 0 for mazes, policies in task.policy_batches(split="train"): - x, order = shuffle(mazes, task.height * task.width) - x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x - output_gpt = reorder(x, order, reverse=True) + output_gpt = eval_mygpt( + gpt, mazes, mode=args.oneshot_input, fixed_len=task.height * task.width + ) output = model(output_gpt) loss = compute_loss(mazes, output, policies, task.height, task.width) @@ -278,9 +281,9 @@ def oneshot(gpt, task): acc_test_loss, nb_test_samples = 0, 0 for mazes, policies in task.policy_batches(split="test"): - x, order = shuffle(mazes, task.height * task.width) - x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x - output_gpt = reorder(x, order, reverse=True) + output_gpt = eval_mygpt( + gpt, mazes, mode=args.oneshot_input, fixed_len=task.height * task.width + ) output = model(output_gpt) loss = compute_loss(mazes, output, policies, task.height, task.width) acc_test_loss += loss.item() * mazes.size(0) @@ -293,9 +296,9 @@ def oneshot(gpt, task): # ------------------- mazes = task.test_input[:32, : task.height * task.width] policies = task.test_policies[:32] - x, order = shuffle(mazes, task.height * task.width) - x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x - output_gpt = reorder(x, order, reverse=True) + output_gpt = eval_mygpt( + gpt, mazes, mode=args.oneshot_input, fixed_len=task.height * task.width + ) output = model(output_gpt) if args.oneshot_output == "policy": targets = policies.permute(0, 2, 1) @@ -628,9 +631,9 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs): for input in task.batches(split="train"): input = input.to(device) - x, order = shuffle(input, task.height * task.width) - x = model(mygpt.BracketedSequence(x), order=order).x - output = reorder(x, order, reverse=True) + output = eval_mygpt( + model, input, mode=args.oneshot_input, fixed_len=task.height * task.width + ) loss = F.cross_entropy(output.transpose(1, 2), input) acc_train_loss += loss.item() * input.size(0) nb_train_samples += input.size(0) -- 2.39.5