X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=beaver.py;h=49cb1f672075795f84021e9e8a5d6573a6d6f56e;hb=10c1e2159ef57a55724fb1753381dc30e8aa77c2;hp=b0fa03c058c416e5d1fad3729e366115d55ba656;hpb=3602bfe2c4e1cd513759bf45cb83f8c2d914674b;p=beaver.git diff --git a/beaver.py b/beaver.py index b0fa03c..49cb1f6 100755 --- a/beaver.py +++ b/beaver.py @@ -26,9 +26,7 @@ else: ###################################################################### -parser = argparse.ArgumentParser( - description="An implementation of GPT with cache to solve a toy geometric reasoning task." -) +parser = argparse.ArgumentParser(description="A maze shortest path solving with a GPT.") parser.add_argument("--log_filename", type=str, default="train.log") @@ -38,9 +36,11 @@ parser.add_argument("--seed", type=int, default=0) parser.add_argument("--nb_epochs", type=int, default=25) -parser.add_argument("--batch_size", type=int, default=100) +parser.add_argument("--nb_train_samples", type=int, default=200000) + +parser.add_argument("--nb_test_samples", type=int, default=50000) -parser.add_argument("--data_size", type=int, default=-1) +parser.add_argument("--batch_size", type=int, default=25) parser.add_argument("--optim", type=str, default="adam") @@ -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) @@ -71,20 +73,27 @@ parser.add_argument("--overwrite_results", action="store_true", default=False) parser.add_argument("--checkpoint_name", type=str, default="checkpoint.pth") ############################## -# picoclvr options +# maze options -parser.add_argument("--world_height", type=int, default=23) +parser.add_argument("--maze_height", type=int, default=13) -parser.add_argument("--world_width", type=int, default=31) +parser.add_argument("--maze_width", type=int, default=21) -parser.add_argument("--world_nb_walls", type=int, default=15) +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() -assert args.prune_properties in {"none", "train+eval", "eval"} - try: os.mkdir(args.result_dir) except FileExistsError: @@ -122,21 +131,39 @@ for n in vars(args): ###################################################################### -def masked_inplace_autoregression( - model, batch_size, input, ar_mask, forbidden_tokens=None, device=torch.device("cpu") -): +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, 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: - model( - mygpt.BracketedSequence(input, 0, i.min()) - ) # Needed to initialize the model's cache + # Needed to initialize the model's cache + 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 forbidden_tokens is not None: - logits = logits.masked_fill(forbidden_tokens, float("-inf")) if args.deterministic_synthesis: t_next = logits.argmax(1) else: @@ -148,8 +175,154 @@ def masked_inplace_autoregression( ###################################################################### +def compute_perplexity(model, split="train"): + with torch.autograd.no_grad(): + t = model.training + model.eval() + + nb_samples, acc_loss = 0, 0.0 + + for input in task.batches(split=split): + input = input.to(device) + 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) + + model.train(t) + + return math.exp(min(100, acc_loss / nb_samples)) + + +###################################################################### + + +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() + + 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, dim_out), + ).to(device) + + for n_epoch in range(args.nb_epochs): + learning_rate = learning_rate_schedule[n_epoch] + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + + acc_train_loss, nb_train_samples = 0, 0 + 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(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 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(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}" + ) + + # ------------------- + 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) + 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) + maze.save_image( + os.path.join( + args.result_dir, + f"oneshot_{args.oneshot_input}_{args.oneshot_output}_{n_epoch:04d}.png", + ), + mazes=mazes, + score_paths=scores, + score_truth=targets, + ) + # ------------------- + + gpt.train(t) + + +###################################################################### + + class Task: - def batches(self, split="train"): + def batches(self, split="train", nb_to_use=-1, desc=None): pass def vocabulary_size(self): @@ -161,288 +334,144 @@ class Task: ###################################################################### -import picoclvr - - -class TaskPicoCLVR(Task): - - # Make a tensor from a list of strings - def tensorize(self, descr): - token_descr = [s.strip().split(" ") for s in descr] - l = max([len(s) for s in token_descr]) - token_descr = [s + [""] * (l - len(s)) for s in token_descr] - id_descr = [[self.token2id[u] for u in s] for s in token_descr] - return torch.tensor(id_descr, device=self.device) - - # Make a list of strings from a tensor - def detensorize(self, x): - return [" ".join([self.id2token[t.item()] for t in r]) for r in x] - - # trim all the tensors in the tuple z to remove as much token from - # left and right in the first tensor. If z is a tuple, all its - # elements are trimed according to the triming for the first - def trim(self, z, token=""): - n = self.token2id[token] - if type(z) == tuple: - x = z[0] - i = (1 - (F.pad(x, (1, 1), value=n) == n).min(0).values.long()).cumsum(0) - a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min() - return tuple([t[:, a:b] for t in z]) - else: - i = (1 - (F.pad(z, (1, 1), value=n) == n).min(0).values.long()).cumsum(0) - a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min() - return z[:, a:b] - - ###################### - # Not the cleanest part of the code - - # Extract the last image of each sequence, from the last - # included, and set to all the tokens from the beginning of - # that image to the end - def excise_last_image(self, input): - t_img, t_nul = self.token2id[""], self.token2id[""] - nb_img_tokens = self.height * self.width + 1 - - input = input.clone() - t = (input == t_img).long() - tail_masks = (t.cumsum(dim=1) == t.sum(dim=1, keepdim=True)).long() - i = (t * tail_masks).nonzero(as_tuple=True) - j = ( - i[0][:, None], - i[1][:, None] + torch.arange(nb_img_tokens, device=input.device)[None, :], - ) - images = self.trim(input[j]) - input[j] = t_nul - loss_masks = 1 - tail_masks - input, loss_masks = self.trim((input, loss_masks)) - return input, loss_masks, images - - def add_true_image(self, input, images, loss_masks): - t_nul = self.token2id[""] - nb_img_tokens = self.height * self.width + 1 - input = F.pad(input, (0, nb_img_tokens), value=t_nul) - loss_masks = F.pad(loss_masks, (0, nb_img_tokens), value=0) - t = (input == t_nul).long() - i = (t.cumsum(dim=1) == 1).nonzero(as_tuple=True) - j = ( - i[0][:, None], - i[1][:, None] + torch.arange(nb_img_tokens, device=input.device)[None, :], - ) - input[j] = images - loss_masks[j] = 1 - input, loss_masks = self.trim((input, loss_masks)) - return input, loss_masks - - def add_generated_image(self, input, loss_masks, model): - t_img, t_nul = self.token2id[""], self.token2id[""] - nb_img_tokens = self.height * self.width + 1 - - input = F.pad(input, (0, nb_img_tokens), value=t_nul) - loss_masks = F.pad(loss_masks, (0, nb_img_tokens), value=0) - t = (input == t_nul).long() - i = (t.cumsum(dim=1) == 1).nonzero(as_tuple=True) - input[i] = t_img - - j = ( - i[0][:, None], - i[1][:, None] - + 1 - + torch.arange(nb_img_tokens - 1, device=input.device)[None, :], - ) - ar_masks = input.new_zeros(input.size(), dtype=torch.int64) - ar_masks[j] = 1 - forbidden_tokens = ( - torch.arange(self.vocabulary_size(), device=input.device) == t_nul - ) - with torch.autograd.no_grad(): - t = model.training - model.eval() - masked_inplace_autoregression( - model, - self.batch_size, - input, - ar_masks, - forbidden_tokens, - device=self.device, - ) - model.train(t) +import maze - input, loss_masks = self.trim((input, loss_masks)) - return input, loss_masks +class TaskMaze(Task): + def map2seq(self, *m): + return torch.cat([x.flatten(1) for x in m], 1) - ###################### + def seq2map(self, s): + s = s.reshape(s.size(0), -1, self.height, self.width) + return (s[:, k] for k in range(s.size(1))) def __init__( self, + nb_train_samples, + nb_test_samples, batch_size, height, width, - nb_colors=5, + nb_walls, device=torch.device("cpu"), - pruner_train=None, - pruner_eval=None, ): - def generate_descr(nb, cache_suffix, pruner): - return picoclvr.generate( - nb, - height=self.height, - width=self.width, - nb_colors=nb_colors, - pruner=pruner, - ) - + self.batch_size = batch_size self.height = height self.width = width - self.batch_size = batch_size self.device = device - nb = args.data_size if args.data_size > 0 else 250000 - self.pruner_train = pruner_train - self.pruner_eval = pruner_eval - - param = { - "nb": nb, - "height": height, - "width": width, - "nb_colors": nb_colors, - "batch_size": batch_size, - "rng_state": list(torch.get_rng_state()), - } - - log_string(f"generating {nb} samples (can take some time)") - self.train_descr = generate_descr( - (nb * 4) // 5, "train", pruner=self.pruner_train + + train_mazes, train_paths, train_policies = maze.create_maze_data( + nb_train_samples, + height=height, + width=width, + nb_walls=nb_walls, + progress_bar=lambda x: tqdm.tqdm(x, dynamic_ncols=True, desc=f"data-train"), ) - self.test_descr = generate_descr((nb * 1) // 5, "test", pruner=None) - - # Build the tokenizer - tokens = {"", ""} - for d in [self.train_descr, self.test_descr]: - for s in d: - for t in s.strip().split(" "): - tokens.add(t) - # make this set a sorted list to get the same tensors given - # the same descr - tokens = list(tokens) - tokens.sort() - self.token2id = dict([(t, n) for n, t in enumerate(tokens)]) - self.id2token = dict([(n, t) for n, t in enumerate(tokens)]) - - # Tokenize the train and test sets - self.train_input = self.tensorize(self.train_descr) - self.test_input = self.tensorize(self.test_descr) - - def batches(self, split="train"): + self.train_input = self.map2seq(train_mazes.to(device), train_paths.to(device)) + self.train_policies = train_policies.flatten(-2).to(device) + + test_mazes, test_paths, test_policies = maze.create_maze_data( + nb_test_samples, + height=height, + width=width, + nb_walls=nb_walls, + 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).to(device) + + self.nb_codes = self.train_input.max() + 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 self.trim(batch) - - def vocabulary_size(self): - return len(self.token2id) + yield batch - def compute_missing_properties(self, n_epoch, model, pruner=None): + 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 + input = input[:, : self.height * self.width] + policies = policies * (input != maze.v_wall)[:, None] - acc_nb_requested_properties = [] - acc_nb_missing_properties = [] - acc_nb_results = 0 + if nb_to_use > 0: + input = input[:nb_to_use] + policies = policies[:nb_to_use] - for input in tqdm.tqdm( - self.test_input.split(self.batch_size), + 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"test-properties", + desc=desc, ): - tape, loss_masks, _ = self.excise_last_image(input) - tape, loss_masks = self.add_generated_image(tape, loss_masks, model) - result_descr = self.detensorize(tape) - np = picoclvr.nb_properties( - result_descr, - height=self.height, - width=self.width, - pruner=pruner, - ) - nb_requested_properties, _, nb_missing_properties = zip(*np) - acc_nb_requested_properties += nb_requested_properties - acc_nb_missing_properties += nb_missing_properties - acc_nb_results += len(result_descr) - - nb_requested_properties = sum(acc_nb_requested_properties) - nb_missing_properties = sum(acc_nb_missing_properties) + yield batch - prefix = "" if pruner is None else "pruned_" - log_string(f"nb_{prefix}samples {n_epoch} {acc_nb_results}") - log_string( - f"property_{prefix}nb {n_epoch} requested {sum(acc_nb_requested_properties)} missing {sum(acc_nb_missing_properties)}" - ) - log_string( - f"property_{prefix}miss {n_epoch} {100*nb_missing_properties/nb_requested_properties:.02f}%" - ) + def vocabulary_size(self): + return self.nb_codes + + def compute_error(self, model, split="train", nb_to_use=-1): + nb_total, nb_correct = 0, 0 + for input in task.batches(split, nb_to_use): + result = input.clone() + ar_mask = result.new_zeros(result.size()) + ar_mask[:, self.height * self.width :] = 1 + result *= 1 - 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) - ###################################################################### + return nb_total, nb_correct def produce_results(self, n_epoch, model): + with torch.autograd.no_grad(): + t = model.training + model.eval() - self.compute_missing_properties(n_epoch, model) - - if self.pruner_eval is not None: - self.compute_missing_properties(n_epoch, model, self.pruner_eval) - - nb_tokens_to_generate = self.height * self.width + 3 - result_descr = [] - nb_per_primer = 8 - primer = [] - - for primer_descr in [ - "red above green green top blue right of red", - "there is red there is yellow there is blue", - "red below yellow yellow below green green below blue red right yellow left green right blue left", - "green bottom yellow bottom green left of blue yellow right of blue blue top", - ]: - primer += [primer_descr] * nb_per_primer - - tape = self.tensorize(primer) - loss_masks = 1 - (tape == self.token2id[""]).long() - tape, loss_masks = self.add_generated_image(tape, loss_masks, model) - result_descr = self.detensorize(tape) - - np = picoclvr.nb_properties(result_descr, height=self.height, width=self.width) - - acc_nb_requested_properties, _, acc_nb_missing_properties = zip(*np) - acc_nb_results = len(result_descr) - - nb_requested_properties = sum(acc_nb_requested_properties) - nb_missing_properties = sum(acc_nb_missing_properties) + train_nb_total, train_nb_correct = self.compute_error( + model, "train", nb_to_use=1000 + ) + log_string( + f"accuracy_train nb_total {train_nb_total} nb_correct {train_nb_correct} accuracy {(100.0*train_nb_correct)/train_nb_total:.02f}%" + ) - prefix = "demo_" - log_string(f"nb_{prefix}samples {n_epoch} {acc_nb_results}") - log_string( - f"property_{prefix}nb {n_epoch} requested {sum(acc_nb_requested_properties)} missing {sum(acc_nb_missing_properties)}" - ) - log_string( - f"property_{prefix}miss {n_epoch} {100*nb_missing_properties/nb_requested_properties:.02f}%" - ) + test_nb_total, test_nb_correct = self.compute_error( + model, "test", nb_to_use=1000 + ) + log_string( + f"accuracy_test nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%" + ) - img = picoclvr.descr2img(result_descr, height=self.height, width=self.width) + input = self.test_input[:32] + result = input.clone() + 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) + + mazes, paths = self.seq2map(input) + _, predicted_paths = self.seq2map(result) + maze.save_image( + os.path.join(args.result_dir, f"result_{n_epoch:04d}.png"), + mazes=mazes, + target_paths=paths, + predicted_paths=predicted_paths, + path_correct=maze.path_correctness(mazes, predicted_paths), + ) - if img.dim() == 5: - if img.size(1) == 1: - img = F.pad(img.squeeze(1), pad=(1, 1, 1, 1), value=64) - else: - img = torch.cat( - [ - torchvision.utils.make_grid(x, padding=1, pad_value=64)[None] - for x in img - ], - 0, - ) - - image_name = os.path.join(args.result_dir, f"result_{n_epoch:04d}.png") - torchvision.utils.save_image( - img / 255.0, image_name, nrow=nb_per_primer, padding=1, pad_value=1.0 - ) - log_string(f"wrote {image_name}") + model.train(t) ###################################################################### @@ -450,24 +479,17 @@ class TaskPicoCLVR(Task): log_string(f"device {device}") -def pruner_horizontal_green(p): - return not ("green" in p and ("left" in p or "right" in p)) - - -task = TaskPicoCLVR( +task = TaskMaze( + nb_train_samples=args.nb_train_samples, + nb_test_samples=args.nb_test_samples, batch_size=args.batch_size, - height=args.height, - width=args.width, - nb_colors=args.nb_colors, + height=args.maze_height, + width=args.maze_width, + nb_walls=args.maze_nb_walls, device=device, - pruner_train=pruner_horizontal_green - if args.prune_properties in {"train+eval"} - else None, - pruner_eval=(lambda p: not pruner_horizontal_green(p)) - if args.prune_properties in {"train+eval", "eval"} - else None, ) + vocabulary_size = task.vocabulary_size() log_string(f"vocabulary_size {vocabulary_size}") @@ -518,8 +540,6 @@ else: ###################################################################### -nb_epochs = args.nb_epochs if args.nb_epochs > 0 else nb_epochs_default - token_count = 0 for input in task.batches(split="train"): token_count += F.one_hot(input, num_classes=task.vocabulary_size()).sum((0, 1)) @@ -553,13 +573,20 @@ log_string(f"learning_rate_schedule {learning_rate_schedule}") ############################## -nb_samples_seen = 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") + + log_string( + f"perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}" + ) -if nb_epochs_finished >= nb_epochs: - task.produce_results(nb_epochs_finished, model) + task.produce_results(n_epoch, model) -for n_epoch in range(nb_epochs_finished, nb_epochs): +############################## +for n_epoch in range(nb_epochs_finished, args.nb_epochs): learning_rate = learning_rate_schedule[n_epoch] log_string(f"learning_rate {learning_rate}") @@ -571,7 +598,7 @@ for n_epoch in range(nb_epochs_finished, 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() @@ -579,41 +606,25 @@ for n_epoch in range(nb_epochs_finished, 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) - nb_samples_seen += input.size(0) optimizer.zero_grad() loss.backward() optimizer.step() - with torch.autograd.no_grad(): - - model.eval() - - nb_test_samples, acc_test_loss = 0, 0.0 - - for input in task.batches(split="test"): - input = input.to(device) - - # input, loss_masks, true_images = task.excise_last_image(input) - # input, loss_masks = task.add_true_image(input, true_images, loss_masks) + train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples)) + test_perplexity = compute_perplexity(model, split="test") - output = model(mygpt.BracketedSequence(input)).x - loss = F.cross_entropy(output.transpose(1, 2), input) - acc_test_loss += loss.item() * input.size(0) - nb_test_samples += input.size(0) - - train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples)) - test_perplexity = math.exp(min(100, acc_test_loss / nb_test_samples)) - - log_string( - f"perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}" - ) + log_string( + f"perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}" + ) - task.produce_results(n_epoch, model) + task.produce_results(n_epoch, model) checkpoint = { "nb_epochs_finished": n_epoch + 1, @@ -629,3 +640,8 @@ for n_epoch in range(nb_epochs_finished, nb_epochs): log_string(f"saved checkpoint {checkpoint_name}") ###################################################################### + +if args.oneshot: + oneshot(model, task) + +######################################################################