From 917e01027ce0c289384dbb87e19d162c695c320b Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Mon, 27 Mar 2023 09:35:29 +0200 Subject: [PATCH] Update --- beaver.py | 36 ++++++++++++++++++------------------ mygpt.py | 7 ++++--- 2 files changed, 22 insertions(+), 21 deletions(-) diff --git a/beaver.py b/beaver.py index 065cda0..5407859 100755 --- a/beaver.py +++ b/beaver.py @@ -133,10 +133,10 @@ for n in vars(args): ###################################################################### -def generation_order(x, fixed_len=0): +def generation_order(x, prompt_len=0): if args.random_regression_order: order = torch.rand(x.size(), device=x.device) - order[:, :fixed_len] = torch.arange(-fixed_len, 0, device=x.device) + order[:, :prompt_len] = torch.arange(-prompt_len, 0, device=x.device) order = order.sort(1).indices else: order = ( @@ -156,13 +156,13 @@ def reorder(x, order, reverse=False): # x is NxTxD1x...xDk, order is NxT' return v -def shuffle(x, fixed_len): - order = generation_order(x, fixed_len) +def shuffle(x, prompt_len): + order = generation_order(x, prompt_len) return reorder(x, order), order -def eval_mygpt(model, input, mode="standard", fixed_len=0): - x, order = shuffle(input, fixed_len) +def eval_mygpt(model, input, mode="standard", prompt_len=0): + x, order = shuffle(input, prompt_len) x = model(mygpt.BracketedSequence(x), mode=mode, order=order).x return reorder(x, order, reverse=True) @@ -195,7 +195,7 @@ def masked_inplace_autoregression(model, batch_size, input, ar_mask, order=None) ###################################################################### -def compute_perplexity(model, task, fixed_len, split="train"): +def compute_perplexity(model, task, prompt_len, split="train"): with torch.autograd.no_grad(): t = model.training model.eval() @@ -204,7 +204,7 @@ def compute_perplexity(model, task, fixed_len, split="train"): for input in task.batches(split=split): input = input.to(device) - output = eval_mygpt(model, input, fixed_len=fixed_len) + output = eval_mygpt(model, input, prompt_len=prompt_len) if args.noncausal_prompt: d = input.size(1) // 2 loss = F.cross_entropy(output[:, d:].transpose(1, 2), input[:, d:]) @@ -273,7 +273,7 @@ def oneshot(gpt, task): acc_train_loss, nb_train_samples = 0, 0 for mazes, policies in task.policy_batches(split="train"): output_gpt = eval_mygpt( - gpt, mazes, mode=args.oneshot_input, fixed_len=task.height * task.width + gpt, mazes, mode=args.oneshot_input, prompt_len=task.height * task.width ) output = model(output_gpt) @@ -288,7 +288,7 @@ def oneshot(gpt, task): acc_test_loss, nb_test_samples = 0, 0 for mazes, policies in task.policy_batches(split="test"): output_gpt = eval_mygpt( - gpt, mazes, mode=args.oneshot_input, fixed_len=task.height * task.width + gpt, mazes, mode=args.oneshot_input, prompt_len=task.height * task.width ) output = model(output_gpt) loss = compute_loss(mazes, output, policies, task.height, task.width) @@ -303,7 +303,7 @@ def oneshot(gpt, task): mazes = task.test_input[:32, : task.height * task.width] policies = task.test_policies[:32] output_gpt = eval_mygpt( - gpt, mazes, mode=args.oneshot_input, fixed_len=task.height * task.width + gpt, mazes, mode=args.oneshot_input, prompt_len=task.height * task.width ) output = model(output_gpt) if args.oneshot_output == "policy": @@ -523,13 +523,15 @@ log_string(f"vocabulary_size {vocabulary_size}") ############################## + def noncausal_prompt_amm_generator(d): q = torch.arange(d)[:, None] k = torch.arange(d)[None, :] s = args.maze_height * args.maze_width -# return torch.logical_and(q < k, torch.logical_or(q >= s, k >= s)) + # return torch.logical_and(q < k, torch.logical_or(q >= s, k >= s)) return q < k + amm_generator = None if args.noncausal_prompt: @@ -616,10 +618,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, task, fixed_len=task.height * task.width, split="train" + model, task, prompt_len=task.height * task.width, split="train" ) test_perplexity = compute_perplexity( - model, task, fixed_len=task.height * task.width, split="test" + model, task, prompt_len=task.height * task.width, split="test" ) log_string( @@ -650,9 +652,7 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs): for input in task.batches(split="train"): input = input.to(device) - output = eval_mygpt( - model, input, fixed_len=task.height * task.width - ) + output = eval_mygpt(model, input, prompt_len=task.height * task.width) if args.noncausal_prompt: d = input.size(1) // 2 loss = F.cross_entropy(output[:, d:].transpose(1, 2), input[:, d:]) @@ -667,7 +667,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, task, fixed_len=task.height * task.width, split="test" + model, task, prompt_len=task.height * task.width, split="test" ) log_string( diff --git a/mygpt.py b/mygpt.py index 06b56df..4555b1e 100755 --- a/mygpt.py +++ b/mygpt.py @@ -148,8 +148,7 @@ class QKVAttention(nn.Module): if amm_generator is None: self.amm_generator = ( - lambda d: torch.arange(d)[:, None] - < torch.arange(d)[None, :] + lambda d: torch.arange(d)[:, None] < torch.arange(d)[None, :] ) else: self.amm_generator = amm_generator @@ -190,7 +189,9 @@ class QKVAttention(nn.Module): if self.causal: if bs_q.first == 0: - self.cache_attzero = self.amm_generator(x_q.size(1)).to(q.device)[None, None,:,:] + self.cache_attzero = self.amm_generator(x_q.size(1)).to(q.device)[ + None, None, :, : + ] a = a.masked_fill( self.cache_attzero[ :, :, bs_q.first : bs_q.first + bs_q.nb, : bs_q.first + bs_q.nb -- 2.39.5