X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=beaver.py;h=49cb1f672075795f84021e9e8a5d6573a6d6f56e;hb=10c1e2159ef57a55724fb1753381dc30e8aa77c2;hp=b505156b49abb389f90f0640fae3d4d59dc9a4e4;hpb=126857a5ef0f205a1d77f62aaf1ee283061396d8;p=beaver.git diff --git a/beaver.py b/beaver.py index b505156..49cb1f6 100755 --- a/beaver.py +++ b/beaver.py @@ -64,6 +64,8 @@ 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("--no_checkpoint", action="store_true", default=False) parser.add_argument("--overwrite_results", action="store_true", default=False) @@ -79,11 +81,14 @@ 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="policy") +parser.add_argument("--oneshot_output", type=str, default="trace") ###################################################################### @@ -126,18 +131,38 @@ for n in vars(args): ###################################################################### +def random_order(result, fixed_len): + if args.random_regression_order: + order = torch.rand(result.size(), device=result.device) + order[:, :fixed_len] = torch.linspace(-2, -1, fixed_len, device=order.device) + return order.sort(1).indices + else: + return torch.arange(result.size(1)).unsqueeze(0).expand(result.size(0), -1) + + +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 + else: + return x.gather(1, order) + + # 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): +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)): 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) @@ -159,8 +184,9 @@ def compute_perplexity(model, split="train"): for input in task.batches(split=split): input = input.to(device) - - output = model(mygpt.BracketedSequence(input)).x + order = random_order(input, task.height * task.width) + input = shuffle(input, order) + output = model(mygpt.BracketedSequence(input), order=order).x loss = F.cross_entropy(output.transpose(1, 2), input) acc_loss += loss.item() * input.size(0) nb_samples += input.size(0) @@ -173,13 +199,21 @@ def compute_perplexity(model, split="train"): ###################################################################### -def oneshot_policy_loss(output, policies, mask): - targets = policies.permute(0, 2, 1) * mask.unsqueeze(-1) - output = output * mask.unsqueeze(-1) - return -(output.log_softmax(-1) * targets).sum() / mask.sum() +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() -# loss = (output.softmax(-1) - targets).abs().max(-1).values.mean() +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): @@ -198,6 +232,7 @@ def oneshot(gpt, task): 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=}") @@ -206,7 +241,7 @@ def oneshot(gpt, task): 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): @@ -214,54 +249,69 @@ def oneshot(gpt, task): optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) acc_train_loss, nb_train_samples = 0, 0 - for input, policies in task.policy_batches(split="train"): - #### - # print(f'{input.size()=} {policies.size()=}') - # s = maze.stationary_densities( - # exit(0) - #### - mask = input == maze.v_empty - output_gpt = gpt(mygpt.BracketedSequence(input), mode=args.oneshot_input).x + for mazes, policies in task.policy_batches(split="train"): + order = random_order(mazes, task.height * task.width) + x = shuffle(mazes, order) + x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x + output_gpt = shuffle(x, order, reorder=True) output = model(output_gpt) - loss = compute_loss(output, policies, mask) - acc_train_loss += loss.item() * input.size(0) - nb_train_samples += input.size(0) + 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, policies in task.policy_batches(split="test"): - mask = input == maze.v_empty - output_gpt = gpt(mygpt.BracketedSequence(input), mode=args.oneshot_input).x + for mazes, policies in task.policy_batches(split="test"): + order = random_order(mazes, task.height * task.width) + x = shuffle(mazes, order) + x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x + output_gpt = shuffle(x, order, reorder=True) output = model(output_gpt) - loss = compute_loss(output, policies, mask) - acc_test_loss += loss.item() * input.size(0) - nb_test_samples += input.size(0) + 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].permute(0, 2, 1) - output_gpt = gpt(mygpt.BracketedSequence(input), mode=args.oneshot_input).x + mazes = task.test_input[:32, : task.height * task.width] + policies = task.test_policies[:32] + order = random_order(mazes, task.height * task.width) + x = shuffle(mazes, order) + x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x + output_gpt = shuffle(x, order, reorder=True) output = model(output_gpt) - scores = ( - (F.one_hot(output.argmax(-1), num_classes=4) * targets).sum(-1) == 0 - ).float() + 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) - input = input.reshape(-1, task.height, task.width) + mazes = mazes.reshape(-1, task.height, task.width) + targets = targets.reshape(-1, task.height, task.width) maze.save_image( os.path.join( args.result_dir, f"oneshot_{args.oneshot_input}_{args.oneshot_output}_{n_epoch:04d}.png", ), - mazes=input, + mazes=mazes, score_paths=scores, + score_truth=targets, ) # ------------------- @@ -272,7 +322,7 @@ def oneshot(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): @@ -332,17 +382,19 @@ class TaskMaze(Task): 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 policies = self.train_policies if split == "train" else self.test_policies @@ -353,10 +405,12 @@ class TaskMaze(Task): input = input[: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), policies.split(self.batch_size)), dynamic_ncols=True, - desc=f"epoch-{split}", + desc=desc, ): yield batch @@ -370,7 +424,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) + order = random_order(result, self.height * self.width) + masked_inplace_autoregression( + model, self.batch_size, result, ar_mask, order=order + ) + result = shuffle(result, order, reorder=True) mazes, paths = self.seq2map(result) nb_correct += maze.path_correctness(mazes, paths).long().sum() nb_total += mazes.size(0) @@ -515,12 +573,6 @@ log_string(f"learning_rate_schedule {learning_rate_schedule}") ############################## -if args.oneshot: - oneshot(model, task) - exit(0) - -############################## - if nb_epochs_finished >= args.nb_epochs: n_epoch = nb_epochs_finished train_perplexity = compute_perplexity(model, split="train") @@ -532,8 +584,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): @@ -556,7 +606,9 @@ 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 + order = random_order(input, task.height * task.width) + input = shuffle(input, order) + 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) nb_train_samples += input.size(0) @@ -588,3 +640,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) + +######################################################################