X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=main.py;h=45bddb762721c468d8990de9cc88044c5cbbd635;hb=f23843d33a4fa5a38f5034deab8f473793732ee3;hp=e741094c1543483514aae45e391d1d3573030a2b;hpb=8f806c62be589d2837f88ffca084ed2ae833124c;p=picoclvr.git diff --git a/main.py b/main.py index e741094..45bddb7 100755 --- a/main.py +++ b/main.py @@ -8,7 +8,7 @@ # torch.backends.cuda.matmul.allow_tf23 # torch.autocast(torch.bfloat16) -import math, sys, argparse, time, tqdm, itertools, os +import math, sys, argparse, time, tqdm, os import torch, torchvision from torch import nn @@ -27,28 +27,33 @@ else: ###################################################################### parser = argparse.ArgumentParser( - description="An implementation of GPT with cache to solve a toy geometric reasoning task." + description="An implementation of GPT with cache.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) -parser.add_argument("--log_filename", type=str, default="train.log") +parser.add_argument( + "--task", type=str, default="picoclvr", help="picoclvr, mnist, maze, snake" +) + +parser.add_argument("--log_filename", type=str, default="train.log", help=" ") parser.add_argument("--result_dir", type=str, default="results_default") parser.add_argument("--seed", type=int, default=0) -parser.add_argument("--nb_epochs", type=int, default=25) +parser.add_argument("--nb_epochs", type=int, default=None) -parser.add_argument("--batch_size", type=int, default=100) +parser.add_argument("--batch_size", type=int, default=None) -parser.add_argument("--data_size", type=int, default=-1) +parser.add_argument("--nb_train_samples", type=int, default=250000) + +parser.add_argument("--nb_test_samples", type=int, default=10000) parser.add_argument("--optim", type=str, default="adam") -parser.add_argument("--learning_rate", type=float, default=1e-3) +parser.add_argument("--learning_rate", type=float, default=1e-4) -parser.add_argument( - "--learning_rate_schedule", type=str, default="10: 2e-4,20: 4e-5,30: 8e-6" -) +parser.add_argument("--learning_rate_schedule", type=str, default="10: 2e-5,30: 4e-6") parser.add_argument("--dim_model", type=int, default=512) @@ -73,19 +78,39 @@ parser.add_argument("--checkpoint_name", type=str, default="checkpoint.pth") ############################## # picoclvr options -parser.add_argument("--nb_colors", type=int, default=5) +parser.add_argument("--picoclvr_nb_colors", type=int, default=5) -parser.add_argument("--height", type=int, default=12) +parser.add_argument("--picoclvr_height", type=int, default=12) -parser.add_argument("--width", type=int, default=16) +parser.add_argument("--picoclvr_width", type=int, default=16) -parser.add_argument("--prune_properties", type=str, default="none") +parser.add_argument("--picocvlr_prune_properties", type=str, default="none") + +############################## +# Maze options + +parser.add_argument("--maze_height", type=int, default=13) + +parser.add_argument("--maze_width", type=int, default=21) + +parser.add_argument("--maze_nb_walls", type=int, default=15) + +############################## +# Snake options + +parser.add_argument("--snake_height", type=int, default=6) + +parser.add_argument("--snake_width", type=int, default=8) + +parser.add_argument("--snake_nb_colors", type=int, default=5) + +parser.add_argument("--snake_length", type=int, default=200) ###################################################################### args = parser.parse_args() -assert args.prune_properties in {"none", "train+eval", "eval"} +assert args.picocvlr_prune_properties in {"none", "train+eval", "eval"} try: os.mkdir(args.result_dir) @@ -106,6 +131,32 @@ if args.seed >= 0: ###################################################################### +default_args = { + "picoclvr": { + "nb_epochs": 25, + "batch_size": 25, + }, + "mnist": { + "nb_epochs": 25, + "batch_size": 10, + }, + "maze": { + "nb_epochs": 25, + "batch_size": 25, + }, + "snake": { + "nb_epochs": 5, + "batch_size": 25, + }, +} + +if args.task in default_args: + for k, v in default_args[args.task].items(): + if getattr(args, k) is None: + setattr(args, k, v) + +###################################################################### + def log_string(s): t = time.strftime("%Y%m%d-%H:%M:%S ", time.localtime()) @@ -124,11 +175,29 @@ for n in vars(args): ###################################################################### +# ra_mask is boolean, with 1s on the values to generate + + def masked_inplace_autoregression( - model, batch_size, input, ar_mask, forbidden_tokens=None, device=torch.device("cpu") + model, + batch_size, + input, + ar_mask, + forbidden_tokens=None, + progress_bar_desc="autoregression", + device=torch.device("cpu"), ): - - for input, ar_mask in zip(input.split(batch_size), ar_mask.split(batch_size)): + # p = logits.softmax(1) + # entropy[:,s]= p.xlogy(p).sum(1) / math.log(2) + batches = zip(input.split(batch_size), ar_mask.split(batch_size)) + if progress_bar_desc is not None: + tqdm.tqdm( + batches, + dynamic_ncols=True, + desc=progress_bar_desc, + total=input.size(0) // batch_size, + ) + for input, ar_mask in batches: i = (ar_mask.sum(0) > 0).nonzero() if i.min() > 0: model( @@ -167,7 +236,6 @@ 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] @@ -265,6 +333,7 @@ class TaskPicoCLVR(Task): input, ar_masks, forbidden_tokens, + progress_bar_desc=None, device=self.device, ) model.train(t) @@ -277,6 +346,8 @@ class TaskPicoCLVR(Task): def __init__( self, + nb_train_samples, + nb_test_samples, batch_size, height, width, @@ -298,12 +369,12 @@ class TaskPicoCLVR(Task): 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, + "nb_train_samples": nb_train_samples, + "nb_test_samples": nb_test_samples, "height": height, "width": width, "nb_colors": nb_colors, @@ -311,11 +382,13 @@ class TaskPicoCLVR(Task): "rng_state": list(torch.get_rng_state()), } - log_string(f"generating {nb} samples (can take some time)") + log_string( + f"generating {nb_train_samples+nb_test_samples} samples (can take some time)" + ) self.train_descr = generate_descr( - (nb * 4) // 5, "train", pruner=self.pruner_train + nb_train_samples, "train", pruner=self.pruner_train ) - self.test_descr = generate_descr((nb * 1) // 5, "test", pruner=None) + self.test_descr = generate_descr(nb_test_samples, "test", pruner=None) # Build the tokenizer tokens = {"", ""} @@ -346,7 +419,6 @@ class TaskPicoCLVR(Task): return len(self.token2id) def compute_missing_properties(self, n_epoch, model, pruner=None): - acc_nb_requested_properties = [] acc_nb_missing_properties = [] acc_nb_results = 0 @@ -385,7 +457,6 @@ class TaskPicoCLVR(Task): ###################################################################### def produce_results(self, n_epoch, model): - self.compute_missing_properties(n_epoch, model) if self.pruner_eval is not None: @@ -440,36 +511,404 @@ class TaskPicoCLVR(Task): 0, ) - image_name = os.path.join(args.result_dir, f"result_{n_epoch:04d}.png") + image_name = os.path.join(args.result_dir, f"picoclvr_result_{n_epoch:04d}.png") torchvision.utils.save_image( - img / 255.0, image_name, nrow=nb_per_primer, padding=1, pad_value=1.0 + img / 255.0, image_name, nrow=nb_per_primer, padding=1, pad_value=0.0 ) log_string(f"wrote {image_name}") ###################################################################### -log_string(f"device {device}") +class TaskMNIST(Task): + def __init__(self, batch_size, device=torch.device("cpu")): + self.device = device + self.batch_size = batch_size + + def batches(self, split="train"): + assert split in {"train", "test"} + data_set = torchvision.datasets.MNIST( + root="./data", train=(split == "train"), download=True + ) + data_input = data_set.data.view(-1, 28 * 28).long() + if args.nb_train_samples is not None: + data_input = data_input[: args.nb_train_samples] + for batch in tqdm.tqdm( + data_input.split(self.batch_size), desc=f"epoch-{split}" + ): + yield batch + + def vocabulary_size(self): + return 256 + + def produce_results(self, n_epoch, model): + results = torch.empty(64, 28 * 28, device=self.device, dtype=torch.int64) + ar_mask = torch.full_like(results, 1) + masked_inplace_autoregression( + model, self.batch_size, results, ar_mask, device=self.device + ) + image_name = os.path.join(args.result_dir, f"mnist_result_{n_epoch:04d}.png") + torchvision.utils.save_image( + 1 - results.reshape(-1, 1, 28, 28) / 255.0, + image_name, + nrow=16, + pad_value=0.8, + ) + log_string(f"wrote {image_name}") + + +###################################################################### + +import maze + + +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_walls, + device=torch.device("cpu"), + ): + self.batch_size = batch_size + self.height = height + self.width = width + self.device = device + + train_mazes, train_paths, _ = 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.train_input = self.map2seq(train_mazes.to(device), train_paths.to(device)) + + test_mazes, test_paths, _ = 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.nb_codes = max(self.train_input.max(), self.test_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=desc + ): + yield batch + + 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 + count = torch.zeros( + self.width * self.height, + self.width * self.height, + device=self.device, + dtype=torch.int64, + ) + for input in tqdm.tqdm( + task.batches(split, nb_to_use), + dynamic_ncols=True, + desc=f"test-mazes", + ): + 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, + progress_bar_desc=None, + device=self.device, + ) + mazes, paths = self.seq2map(result) + path_correctness = maze.path_correctness(mazes, paths) + nb_correct += path_correctness.long().sum() + nb_total += mazes.size(0) + + optimal_path_lengths = ( + (input[:, self.height * self.width :] == maze.v_path).long().sum(1) + ) + predicted_path_lengths = ( + (result[:, self.height * self.width :] == maze.v_path).long().sum(1) + ) + optimal_path_lengths = optimal_path_lengths[path_correctness] + predicted_path_lengths = predicted_path_lengths[path_correctness] + count[optimal_path_lengths, predicted_path_lengths] += 1 + + if count.max() == 0: + count = None + else: + count = count[ + : count.sum(1).nonzero().max() + 1, : count.sum(0).nonzero().max() + 1 + ] + + return nb_total, nb_correct, count + + def produce_results(self, n_epoch, model): + with torch.autograd.no_grad(): + t = model.training + model.eval() + + train_nb_total, train_nb_correct, count = 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}%" + ) -def pruner_horizontal_green(p): + test_nb_total, test_nb_correct, count = 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}%" + ) + + if count is not None: + proportion_optimal = count.diagonal().sum().float() / count.sum() + log_string(f"proportion_optimal_test {proportion_optimal*100:.02f}%") + with open( + os.path.join(args.result_dir, f"maze_result_{n_epoch:04d}.txt"), "w" + ) as f: + for i in range(count.size(0)): + for j in range(count.size(1)): + eol = " " if j < count.size(1) - 1 else "\n" + f.write(f"{count[i,j]}{eol}") + + input = self.test_input[:48] + 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, device=self.device + ) + + mazes, paths = self.seq2map(input) + _, predicted_paths = self.seq2map(result) + + filename = os.path.join(args.result_dir, f"maze_result_{n_epoch:04d}.png") + maze.save_image( + filename, + mazes=mazes, + target_paths=paths, + predicted_paths=predicted_paths, + path_correct=maze.path_correctness(mazes, predicted_paths), + path_optimal=maze.path_optimality(paths, predicted_paths), + ) + log_string(f"wrote {filename}") + + model.train(t) + + +###################################################################### + + +import snake + + +class TaskSnake(Task): + def __init__( + self, + nb_train_samples, + nb_test_samples, + batch_size, + height, + width, + nb_colors, + length, + prompt_length, + device=torch.device("cpu"), + ): + self.batch_size = batch_size + self.height = height + self.width = width + self.device = device + self.prompt_length = prompt_length + + self.train_input, self.train_prior_visits, _, _ = snake.generate_sequences( + nb_train_samples, + height, + width, + nb_colors, + length, + prompt_length, + self.device, + ) + self.test_input, self.test_prior_visits, _, _ = snake.generate_sequences( + nb_test_samples, + height, + width, + nb_colors, + length, + prompt_length, + self.device, + ) + + self.nb_codes = max(self.train_input.max(), self.test_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=desc + ): + yield batch + + def vocabulary_size(self): + return self.nb_codes + + def produce_results(self, n_epoch, model): + with torch.autograd.no_grad(): + t = model.training + model.eval() + + def compute_nb_correct(input, prior_visits): + result = input.clone() + i = torch.arange(result.size(1), device=result.device)[None, :] + ar_mask = ( + torch.logical_and(i >= self.prompt_length * 2, i % 2 == 0) + .long() + .expand_as(result) + ) + result *= 1 - ar_mask + + # snake.solver(result,ar_mask) + + masked_inplace_autoregression( + model, self.batch_size, result, ar_mask, device=self.device + ) + + nb_total = ((prior_visits > 0) * ar_mask).sum() + + nb_correct = ( + (result == input).long() * (prior_visits > 0) * ar_mask + ).sum() + + # nb_total = result.size(0) + # nb_correct = ((result - input).abs().sum(1) == 0).sum() + + return nb_total, nb_correct + + # train_nb_total, train_nb_correct = compute_nb_correct( + # self.train_input, self.train_prior_visits + # ) + + # 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}%" + # ) + + test_nb_total, test_nb_correct = compute_nb_correct( + self.test_input[:1000], self.test_prior_visits[: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}%" + ) + + model.train(t) + + +###################################################################### + + +def picoclvr_pruner_horizontal_green(p): return not ("green" in p and ("left" in p or "right" in p)) -task = TaskPicoCLVR( - batch_size=args.batch_size, - height=args.height, - width=args.width, - nb_colors=args.nb_colors, - 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, +picoclvr_pruner_train = ( + picoclvr_pruner_horizontal_green + if args.picocvlr_prune_properties in {"train+eval"} + else None ) +picoclvr_pruner_eval = ( + (lambda p: not picoclvr_pruner_horizontal_green(p)) + if args.picocvlr_prune_properties in {"train+eval", "eval"} + else None +) + +###################################################################### + +if args.task == "picoclvr": + task = TaskPicoCLVR( + nb_train_samples=args.nb_train_samples, + nb_test_samples=args.nb_test_samples, + batch_size=args.batch_size, + height=args.picoclvr_height, + width=args.picoclvr_width, + nb_colors=args.picoclvr_nb_colors, + device=device, + pruner_train=picoclvr_pruner_train, + pruner_eval=picoclvr_pruner_eval, + ) + +elif args.task == "mnist": + task = TaskMNIST( + batch_size=args.batch_size, + device=device, + ) + +elif args.task == "maze": + task = TaskMaze( + nb_train_samples=args.nb_train_samples, + nb_test_samples=args.nb_test_samples, + batch_size=args.batch_size, + height=args.maze_height, + width=args.maze_width, + nb_walls=args.maze_nb_walls, + device=device, + ) + +elif args.task == "snake": + task = TaskSnake( + nb_train_samples=args.nb_train_samples, + nb_test_samples=args.nb_test_samples, + batch_size=args.batch_size, + height=args.snake_height, + width=args.snake_width, + nb_colors=args.snake_nb_colors, + length=args.snake_length, + prompt_length=args.snake_length // 2, + device=device, + ) + +else: + raise ValueError(f"Unknown task {args.task}") + +###################################################################### + +log_string(f"device {device}") + vocabulary_size = task.vocabulary_size() log_string(f"vocabulary_size {vocabulary_size}") @@ -561,7 +1000,6 @@ if nb_epochs_finished >= nb_epochs: task.produce_results(nb_epochs_finished, model) for n_epoch in range(nb_epochs_finished, nb_epochs): - learning_rate = learning_rate_schedule[n_epoch] log_string(f"learning_rate {learning_rate}") @@ -592,7 +1030,6 @@ for n_epoch in range(nb_epochs_finished, nb_epochs): optimizer.step() with torch.autograd.no_grad(): - model.eval() nb_test_samples, acc_test_loss = 0, 0.0 @@ -600,9 +1037,6 @@ for n_epoch in range(nb_epochs_finished, nb_epochs): 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) - output = model(mygpt.BracketedSequence(input)).x loss = F.cross_entropy(output.transpose(1, 2), input) acc_test_loss += loss.item() * input.size(0)