X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=beaver.py;h=f850f699ed9c3d18ef521df99b0a77595485fa7c;hb=e14c55948c1f099f95fa3d7343b5c939e60fcb1c;hp=de38ff40ca373d695413039a745ac14e65d4ee30;hpb=e23374a06a07a1bc899e1c7ff7f5d8be75f9cdb5;p=beaver.git diff --git a/beaver.py b/beaver.py index de38ff4..f850f69 100755 --- a/beaver.py +++ b/beaver.py @@ -64,12 +64,14 @@ parser.add_argument("--dropout", type=float, default=0.1) parser.add_argument("--deterministic_synthesis", action="store_true", default=False) +parser.add_argument("--random_regression_order", action="store_true", default=False) + +parser.add_argument("--noncausal_prompt", action="store_true", default=False) + parser.add_argument("--no_checkpoint", action="store_true", default=False) parser.add_argument("--overwrite_results", action="store_true", default=False) -parser.add_argument("--one_shot", action="store_true", default=False) - parser.add_argument("--checkpoint_name", type=str, default="checkpoint.pth") ############################## @@ -81,6 +83,15 @@ 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") + +parser.add_argument("--oneshot_output", type=str, default="trace") + ###################################################################### args = parser.parse_args() @@ -122,18 +133,56 @@ for n in vars(args): ###################################################################### +def generation_order(x, fixed_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 = order.sort(1).indices + else: + order = ( + torch.arange(x.size(1), device=x.device).unsqueeze(0).expand(x.size(0), -1) + ) + return order + + +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()).scatter_(1, order, u) + else: + 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 + + +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 # tokens that should be generated -def masked_inplace_autoregression(model, batch_size, input, ar_mask): - for input, ar_mask in zip(input.split(batch_size), ar_mask.split(batch_size)): +def masked_inplace_autoregression(model, batch_size, input, ar_mask, order=None): + 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 - model(mygpt.BracketedSequence(input, 0, i.min())) + model(mygpt.BracketedSequence(input, 0, i.min()), order=order) for s in range(i.min(), i.max() + 1): - output = model(mygpt.BracketedSequence(input, s, 1)).x + output = model(mygpt.BracketedSequence(input, s, 1), order=order).x logits = output[:, s] if args.deterministic_synthesis: t_next = logits.argmax(1) @@ -146,7 +195,7 @@ def masked_inplace_autoregression(model, batch_size, input, ar_mask): ###################################################################### -def compute_perplexity(model, split="train"): +def compute_perplexity(model, task, fixed_len, split="train"): with torch.autograd.no_grad(): t = model.training model.eval() @@ -155,9 +204,12 @@ def compute_perplexity(model, split="train"): for input in task.batches(split=split): input = input.to(device) - - output = model(mygpt.BracketedSequence(input)).x - loss = F.cross_entropy(output.transpose(1, 2), input) + output = eval_mygpt(model, input, fixed_len=fixed_len) + if args.noncausal_prompt: + t = input.size(1) // 2 + loss = F.cross_entropy(output[:, t:].transpose(1, 2), input[:, t:]) + else: + loss = F.cross_entropy(output.transpose(1, 2), input) acc_loss += loss.item() * input.size(0) nb_samples += input.size(0) @@ -169,17 +221,49 @@ def compute_perplexity(model, split="train"): ###################################################################### -def one_shot(gpt, task): +def oneshot_policy_loss(mazes, output, policies, height, width): + masks = (mazes == maze.v_empty).unsqueeze(-1) + targets = policies.permute(0, 2, 1) * masks + output = output * masks + return -(output.log_softmax(-1) * targets).sum() / masks.sum() + + +def oneshot_trace_loss(mazes, output, policies, height, width): + masks = mazes == maze.v_empty + targets = maze.stationary_densities( + mazes.view(-1, height, width), policies.view(-1, 4, height, width) + ).flatten(-2) + targets = targets * masks + output = output.squeeze(-1) * masks + return (output - targets).abs().sum() / masks.sum() + + +def oneshot(gpt, task): t = gpt.training gpt.eval() - mode='head' - dim_in=args.dim_model * (args.nb_blocks * 2 if mode=='deep' else 1) + + if args.oneshot_input == "head": + dim_in = args.dim_model + elif args.oneshot_input == "deep": + dim_in = args.dim_model * args.nb_blocks * 2 + else: + raise ValueError(f"{args.oneshot_input=}") + + if args.oneshot_output == "policy": + dim_out = 4 + compute_loss = oneshot_policy_loss + elif args.oneshot_output == "trace": + dim_out = 1 + compute_loss = oneshot_trace_loss + else: + raise ValueError(f"{args.oneshot_output=}") + model = nn.Sequential( nn.Linear(dim_in, args.dim_model), nn.ReLU(), nn.Linear(args.dim_model, args.dim_model), nn.ReLU(), - nn.Linear(args.dim_model, 4), + nn.Linear(args.dim_model, dim_out), ).to(device) for n_epoch in range(args.nb_epochs): @@ -187,61 +271,69 @@ def one_shot(gpt, task): optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) acc_train_loss, nb_train_samples = 0, 0 - for input, targets in task.policy_batches(split="train"): - output_gpt = gpt(mygpt.BracketedSequence(input), mode=mode).x - output = model(output_gpt) - targets = targets * (input.unsqueeze(-1) == maze.v_empty) - output = output * (input.unsqueeze(-1) == maze.v_empty) - # loss = (output.softmax(-1) - targets).abs().max(-1).values.mean() - loss = ( - -(output.log_softmax(-1) * targets).sum() - / (input == maze.v_empty).sum() + 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 ) - acc_train_loss += loss.item() * input.size(0) - nb_train_samples += input.size(0) + output = model(output_gpt) + + loss = compute_loss(mazes, output, policies, task.height, task.width) + acc_train_loss += loss.item() * mazes.size(0) + nb_train_samples += mazes.size(0) optimizer.zero_grad() loss.backward() optimizer.step() acc_test_loss, nb_test_samples = 0, 0 - for input, targets in task.policy_batches(split="test"): - output_gpt = gpt(mygpt.BracketedSequence(input), mode=mode).x - output = model(output_gpt) - targets = targets * (input.unsqueeze(-1) == maze.v_empty) - output = output * (input.unsqueeze(-1) == maze.v_empty) - # loss = (output.softmax(-1) - targets).abs().max(-1).values.mean() - loss = ( - -(output.log_softmax(-1) * targets).sum() - / (input == maze.v_empty).sum() + 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 ) - acc_test_loss += loss.item() * input.size(0) - nb_test_samples += input.size(0) + output = model(output_gpt) + loss = compute_loss(mazes, output, policies, task.height, task.width) + acc_test_loss += loss.item() * mazes.size(0) + nb_test_samples += mazes.size(0) log_string( f"diff_ce {n_epoch} train {acc_train_loss/nb_train_samples} test {acc_test_loss/nb_test_samples}" ) # ------------------- - input = task.test_input[:32, : task.height * task.width] - targets = task.test_policies[:32] - output_gpt = gpt(mygpt.BracketedSequence(input), mode=mode).x + 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 + ) output = model(output_gpt) - # losses = (-output.log_softmax(-1) * targets + targets.xlogy(targets)).sum(-1) - # losses = losses * (input == maze.v_empty) - # losses = losses / losses.max() - # losses = (output.softmax(-1) - targets).abs().max(-1).values - # losses = (losses >= 0.05).float() - losses = ( - (F.one_hot(output.argmax(-1), num_classes=4) * targets).sum(-1) == 0 - ).float() - losses = losses.reshape(-1, args.maze_height, args.maze_width) - input = input.reshape(-1, args.maze_height, args.maze_width) + if args.oneshot_output == "policy": + targets = policies.permute(0, 2, 1) + scores = ( + (F.one_hot(output.argmax(-1), num_classes=4) * targets).sum(-1) == 0 + ).float() + elif args.oneshot_output == "trace": + targets = maze.stationary_densities( + mazes.view(-1, task.height, task.width), + policies.view(-1, 4, task.height, task.width), + ).flatten(-2) + scores = output + else: + raise ValueError(f"{args.oneshot_output=}") + + scores = scores.reshape(-1, task.height, task.width) + mazes = mazes.reshape(-1, task.height, task.width) + targets = targets.reshape(-1, task.height, task.width) + filename = ( + f"oneshot_{args.oneshot_input}_{args.oneshot_output}_{n_epoch:04d}.png" + ) maze.save_image( - os.path.join(args.result_dir, f"oneshot_{n_epoch:04d}.png"), - mazes=input, - score_paths=losses, + os.path.join(args.result_dir, filename), + mazes=mazes, + score_paths=scores, + score_truth=targets, ) + log_string(f"wrote {filename}") + # ------------------- gpt.train(t) @@ -251,7 +343,7 @@ def one_shot(gpt, task): class Task: - def batches(self, split="train"): + def batches(self, split="train", nb_to_use=-1, desc=None): pass def vocabulary_size(self): @@ -297,7 +389,7 @@ class TaskMaze(Task): progress_bar=lambda x: tqdm.tqdm(x, dynamic_ncols=True, desc=f"data-train"), ) self.train_input = self.map2seq(train_mazes.to(device), train_paths.to(device)) - self.train_policies = train_policies.flatten(-2).permute(0, 2, 1).to(device) + self.train_policies = train_policies.flatten(-2).to(device) test_mazes, test_paths, test_policies = maze.create_maze_data( nb_test_samples, @@ -307,35 +399,39 @@ class TaskMaze(Task): progress_bar=lambda x: tqdm.tqdm(x, dynamic_ncols=True, desc=f"data-test"), ) self.test_input = self.map2seq(test_mazes.to(device), test_paths.to(device)) - self.test_policies = test_policies.flatten(-2).permute(0, 2, 1).to(device) + self.test_policies = test_policies.flatten(-2).to(device) self.nb_codes = self.train_input.max() + 1 - def batches(self, split="train", nb_to_use=-1): + def batches(self, split="train", nb_to_use=-1, desc=None): assert split in {"train", "test"} input = self.train_input if split == "train" else self.test_input if nb_to_use > 0: input = input[:nb_to_use] + if desc is None: + desc = f"epoch-{split}" for batch in tqdm.tqdm( - input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}" + input.split(self.batch_size), dynamic_ncols=True, desc=desc ): yield batch - def policy_batches(self, split="train", nb_to_use=-1): + def policy_batches(self, split="train", nb_to_use=-1, desc=None): assert split in {"train", "test"} input = self.train_input if split == "train" else self.test_input - targets = self.train_policies if split == "train" else self.test_policies + policies = self.train_policies if split == "train" else self.test_policies input = input[:, : self.height * self.width] - targets = targets * (input != maze.v_wall)[:, :, None] + policies = policies * (input != maze.v_wall)[:, None] if nb_to_use > 0: input = input[:nb_to_use] - targets = targets[:nb_to_use] + policies = policies[:nb_to_use] + if desc is None: + desc = f"epoch-{split}" for batch in tqdm.tqdm( - zip(input.split(self.batch_size), targets.split(self.batch_size)), + zip(input.split(self.batch_size), policies.split(self.batch_size)), dynamic_ncols=True, - desc=f"epoch-{split}", + desc=desc, ): yield batch @@ -349,7 +445,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, reverse=True) mazes, paths = self.seq2map(result) nb_correct += maze.path_correctness(mazes, paths).long().sum() nb_total += mazes.size(0) @@ -380,17 +480,23 @@ 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, reverse=True) mazes, paths = self.seq2map(input) _, predicted_paths = self.seq2map(result) + filename = f"result_{n_epoch:04d}.png" maze.save_image( - os.path.join(args.result_dir, f"result_{n_epoch:04d}.png"), + os.path.join(args.result_dir, filename), mazes=mazes, target_paths=paths, predicted_paths=predicted_paths, path_correct=maze.path_correctness(mazes, predicted_paths), ) + log_string(f"wrote {filename}") model.train(t) @@ -417,6 +523,17 @@ log_string(f"vocabulary_size {vocabulary_size}") ############################## +amm_generator = None + +if args.noncausal_prompt: + amm_generator = lambda d: torch.logical_and( + torch.arange(d)[None, None, :, None] < torch.arange(d)[None, None, None, :], + torch.logical_or( + torch.arange(d)[None, None, :, None] >= d // 2, + torch.arange(d)[None, None, None, :] >= d // 2, + ), + ) + model = mygpt.MyGPT( vocabulary_size=vocabulary_size, dim_model=args.dim_model, @@ -426,6 +543,7 @@ model = mygpt.MyGPT( nb_blocks=args.nb_blocks, causal=True, dropout=args.dropout, + amm_generator=amm_generator, ) model.to(device) @@ -494,16 +612,14 @@ log_string(f"learning_rate_schedule {learning_rate_schedule}") ############################## -if args.one_shot: - one_shot(model, task) - exit(0) - -############################## - if nb_epochs_finished >= args.nb_epochs: n_epoch = nb_epochs_finished - train_perplexity = compute_perplexity(model, split="train") - test_perplexity = compute_perplexity(model, split="test") + train_perplexity = compute_perplexity( + model, task, fixed_len=task.height * task.width, split="train" + ) + test_perplexity = compute_perplexity( + model, task, fixed_len=task.height * task.width, split="test" + ) log_string( f"perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}" @@ -511,8 +627,6 @@ if nb_epochs_finished >= args.nb_epochs: task.produce_results(n_epoch, model) - exit(0) - ############################## for n_epoch in range(nb_epochs_finished, args.nb_epochs): @@ -527,7 +641,7 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs): elif args.optim == "adamw": optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) else: - raise ValueError(f"Unknown optimizer {args.optim}.") + raise ValueError(f"{args.optim=}") model.train() @@ -535,8 +649,14 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs): for input in task.batches(split="train"): input = input.to(device) - output = model(mygpt.BracketedSequence(input)).x - loss = F.cross_entropy(output.transpose(1, 2), input) + output = eval_mygpt( + model, input, mode=args.oneshot_input, fixed_len=task.height * task.width + ) + if args.noncausal_prompt: + t = input.size(1) // 2 + loss = F.cross_entropy(output[:, t:].transpose(1, 2), input[:, t:]) + else: + loss = F.cross_entropy(output.transpose(1, 2), input) acc_train_loss += loss.item() * input.size(0) nb_train_samples += input.size(0) @@ -545,7 +665,9 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs): optimizer.step() train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples)) - test_perplexity = compute_perplexity(model, split="test") + test_perplexity = compute_perplexity( + model, task, fixed_len=task.height * task.width, split="test" + ) log_string( f"perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}" @@ -567,3 +689,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) + +######################################################################