X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=main.py;h=11eb8fd90053e07d156a53cdc3f57cf62c50562d;hb=66d210bd5e04ae58f9e1495df77f1f975ee99c56;hp=6d9f69d65120dc5152c4a38ed2cf72a7a7682b72;hpb=943a440a83b98de60bad767a9ad09f63b5088514;p=culture.git diff --git a/main.py b/main.py index 6d9f69d..11eb8fd 100755 --- a/main.py +++ b/main.py @@ -5,91 +5,189 @@ # Written by Francois Fleuret -import math, sys, argparse, time, tqdm, itertools, os +import math, sys, argparse, time, tqdm, os, datetime, warnings import torch, torchvision from torch import nn from torch.nn import functional as F -import mygpt, tensorstack +import ffutils +import mygpt +import sky, wireworld, quizz_machine + +# world quizzes vs. culture quizzes ###################################################################### -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +nb_new_c_quizzes_for_train = 1000 +nb_new_c_quizzes_for_test = 100 + +###################################################################### + +if torch.cuda.is_available(): + device = torch.device("cuda") + torch.backends.cuda.matmul.allow_tf32 = True +else: + device = torch.device("cpu") ###################################################################### parser = argparse.ArgumentParser( - description="An implementation of GPT with cache to solve a toy geometric reasonning 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("--result_dir", type=str, default="results_default") +parser.add_argument("--result_dir", type=str, default=None) parser.add_argument("--seed", type=int, default=0) -parser.add_argument("--nb_epochs", type=int, default=25) +parser.add_argument("--max_percents_of_test_in_train", type=int, default=1) -parser.add_argument("--batch_size", type=int, default=100) +######################################## -parser.add_argument("--data_size", type=int, default=-1) +parser.add_argument("--nb_epochs", type=int, default=10000) -parser.add_argument("--optim", type=str, default="adam") +parser.add_argument("--batch_size", type=int, default=None) + +parser.add_argument("--physical_batch_size", type=int, default=None) + +parser.add_argument("--nb_train_samples", type=int, default=None) + +parser.add_argument("--nb_test_samples", type=int, default=None) parser.add_argument("--learning_rate", type=float, default=1e-3) -parser.add_argument( - "--learning_rate_schedule", type=str, default="10: 2e-4,20: 4e-5,30: 8e-6" -) +######################################## -parser.add_argument("--dim_model", type=int, default=512) +parser.add_argument("--model", type=str, default=None) -parser.add_argument("--dim_keys", type=int, default=64) +parser.add_argument("--dim_model", type=int, default=None) -parser.add_argument("--dim_hidden", type=int, default=2048) +parser.add_argument("--dim_keys", type=int, default=None) -parser.add_argument("--nb_heads", type=int, default=8) +parser.add_argument("--dim_hidden", type=int, default=None) -parser.add_argument("--nb_blocks", type=int, default=12) +parser.add_argument("--nb_heads", type=int, default=None) + +parser.add_argument("--nb_blocks", type=int, default=None) parser.add_argument("--dropout", type=float, default=0.1) -parser.add_argument("--nb_oneshot_blocks", type=int, default=-1) +######################################## parser.add_argument("--deterministic_synthesis", action="store_true", default=False) -parser.add_argument("--no_checkpoint", action="store_true", default=False) +parser.add_argument("--problem", type=str, default="sky") -parser.add_argument("--overwrite_results", action="store_true", default=False) +parser.add_argument("--nb_gpts", type=int, default=5) -parser.add_argument("--checkpoint_name", type=str, default="checkpoint.pth") +parser.add_argument("--nb_models_for_generation", type=int, default=1) -############################## -# picoclvr options +parser.add_argument("--generation_mode", type=str, default="groupthink") + +parser.add_argument("--min_to_validate", type=int, default=4) + +parser.add_argument("--max_to_validate", type=int, default=4) -parser.add_argument("--nb_colors", type=int, default=5) +parser.add_argument("--accuracy_to_make_c_quizzes", type=float, default=0.975) -parser.add_argument("--height", type=int, default=12) +parser.add_argument("--dirty_debug", action="store_true", default=False) -parser.add_argument("--width", type=int, default=16) +parser.add_argument("--sky_height", type=int, default=6) -parser.add_argument("--prune_properties", type=str, default="none") +parser.add_argument("--sky_width", type=int, default=8) + +parser.add_argument("--sky_nb_birds", type=int, default=3) + +parser.add_argument("--sky_nb_iterations", type=int, default=2) + +parser.add_argument("--sky_speed", type=int, default=3) ###################################################################### args = parser.parse_args() -assert args.prune_properties in {"none", "train+eval", "eval"} +if args.result_dir is None: + args.result_dir = f"results_culture" + +###################################################################### + +if args.dirty_debug: + args.accuracy_to_make_c_quizzes = 0.0 + nb_new_c_quizzes_for_train = 100 + nb_new_c_quizzes_for_test = 10 + +###################################################################### + +default_args = { + "model": "37M", + "batch_size": 100, + "nb_train_samples": 100000, + "nb_test_samples": 10000, +} + +for k, v in default_args.items(): + if getattr(args, k) is None: + setattr(args, k, v) + +###################################################################### + +default_model_args = { + "17K": { + "dim_model": 32, + "dim_keys": 32, + "dim_hidden": 32, + "nb_heads": 2, + "nb_blocks": 2, + }, + "4M": { + "dim_model": 256, + "dim_keys": 32, + "dim_hidden": 1024, + "nb_heads": 4, + "nb_blocks": 6, + }, + "37M": { + "dim_model": 512, + "dim_keys": 64, + "dim_hidden": 2048, + "nb_heads": 8, + "nb_blocks": 12, + }, + "122M": { + "dim_model": 768, + "dim_keys": 64, + "dim_hidden": 2048, + "nb_heads": 8, + "nb_blocks": 24, + }, + "352M": { + "dim_model": 1024, + "dim_keys": 64, + "dim_hidden": 2048, + "nb_heads": 8, + "nb_blocks": 48, + }, +} + +if args.model in default_model_args: + for k, v in default_model_args[args.model].items(): + if getattr(args, k) is None: + setattr(args, k, v) +else: + raise ValueError(f"Unknown model {args.model}") + +###################################################################### try: os.mkdir(args.result_dir) except FileExistsError: - if not args.overwrite_results: - print(f"result directory {args.result_dir} already exists") - exit(1) + print(f"result directory {args.result_dir} already exists") + exit(1) -log_file = open(os.path.join(args.result_dir, args.log_filename), "w") +log_file = open(os.path.join(args.result_dir, args.log_filename), "a") if args.seed >= 0: # torch.backends.cudnn.deterministic = True @@ -113,518 +211,352 @@ def log_string(s): sys.stdout.flush() +log_string(f"argv {' '.join(sys.argv)}") + for n in vars(args): log_string(f"args.{n} {getattr(args, n)}") -###################################################################### - -def masked_inplace_autoregression( - model, batch_size, input, ar_mask, forbidden_tokens=None, device=torch.device("cpu") -): +###################################################################### - 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 - for s in range(i.min(), i.max() + 1): - output = model(mygpt.BracketedSequence(input, s, 1)).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: - dist = torch.distributions.categorical.Categorical(logits=logits) - t_next = dist.sample() - input[:, s] = ar_mask[:, s] * t_next + (1 - ar_mask[:, s]) * input[:, s] +if args.dirty_debug: + args.nb_train_samples = 2500 + args.nb_test_samples = 100 +if args.physical_batch_size is None: + args.physical_batch_size = args.batch_size +else: + assert args.batch_size % args.physical_batch_size == 0 + +assert args.nb_train_samples % args.batch_size == 0 +assert args.nb_test_samples % args.batch_size == 0 + +if args.problem == "sky": + problem = sky.Sky( + height=args.sky_height, + width=args.sky_width, + nb_birds=args.sky_nb_birds, + nb_iterations=args.sky_nb_iterations, + speed=args.sky_speed, + ) +elif args.problem == "wireworld": + problem = wireworld.Wireworld(height=8, width=10, nb_iterations=2, speed=5) +else: + raise ValueError + +quizz_machine = quizz_machine.QuizzMachine( + problem=problem, + nb_train_samples=args.nb_train_samples, + nb_test_samples=args.nb_test_samples, + batch_size=args.physical_batch_size, + result_dir=args.result_dir, + logger=log_string, + device=device, +) ###################################################################### +log_string(f"device {device}") -class Task: - def batches(self, split="train"): - pass - - def vocabulary_size(self): - pass - - def produce_results(self, n_epoch, model): - pass +vocabulary_size = quizz_machine.vocabulary_size() +log_string(f"vocabulary_size {vocabulary_size}") ###################################################################### -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) - - input, loss_masks = self.trim((input, loss_masks)) - - return input, loss_masks +# Compute the entropy of the training tokens - ###################### +token_count = 0 +for input in quizz_machine.batches(split="train", desc="train-entropy"): + token_count += F.one_hot(input, num_classes=quizz_machine.vocabulary_size()).sum( + (0, 1) + ) +token_probas = token_count / token_count.sum() +entropy = -torch.xlogy(token_probas, token_probas).sum() +train_set_perplexity = math.exp(entropy) - def __init__( - self, - batch_size, - height, - width, - nb_colors=5, - device=torch.device("cpu"), - pruner_train=None, - pruner_eval=None, +###################################################################### +# A bit of paranoia never hurts + +if args.max_percents_of_test_in_train >= 0: + + def subsets_as_tuples(batches, cs): + s = set() + for batch in batches: + for x in batch: + s.add(tuple([v.item() for v in x])) + if len(s) == cs: + yield s + s = set() + yield s + + nb_test, nb_in_train = 0, 0 + for test_subset in subsets_as_tuples( + quizz_machine.batches(split="test", desc="test-check"), 25000 ): - def generate_descr(nb, cache_suffix, pruner): - return picoclvr.generate( - nb, - height=self.height, - width=self.width, - nb_colors=nb_colors, - pruner=pruner, - ) - - 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 - ) - 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"): - assert split in {"train", "test"} - input = self.train_input if split == "train" else self.test_input - for batch in tqdm.tqdm( - input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}" + in_train = set() + for train_subset in subsets_as_tuples( + quizz_machine.batches(split="train", desc="train-check"), 25000 ): - yield self.trim(batch) + in_train.update(test_subset.intersection(train_subset)) + nb_in_train += len(in_train) + nb_test += len(test_subset) - def vocabulary_size(self): - return len(self.token2id) + log_string( + f"data_check {nb_in_train*100/nb_test:.02f}% ({nb_in_train}/{nb_test}) of test samples are in the train set" + ) - def compute_missing_properties(self, n_epoch, model, pruner=None): + assert ( + nb_in_train <= args.max_percents_of_test_in_train * nb_test / 100 + ), f"More than {args.max_percents_of_test_in_train}% of test samples are in the train set" - acc_nb_requested_properties = [] - acc_nb_missing_properties = [] - acc_nb_results = 0 +############################## - for input in tqdm.tqdm( - self.test_input.split(self.batch_size), - dynamic_ncols=True, - desc=f"test-properties", - ): - 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) - - 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 one_epoch(model, quizz_machine): + optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate) - def produce_results(self, n_epoch, model): + model.train() - self.compute_missing_properties(n_epoch, model) + nb_train_samples, acc_train_loss = 0, 0.0 - if self.pruner_eval is not None: - self.compute_missing_properties(n_epoch, model, self.pruner_eval) + for input in quizz_machine.batches(split="train"): + input = input.to(device) - nb_tokens_to_generate = self.height * self.width + 3 - result_descr = [] - nb_per_primer = 8 - primer = [] + if nb_train_samples % args.batch_size == 0: + optimizer.zero_grad() - 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 + output = model(mygpt.BracketedSequence(input)).x + loss = F.cross_entropy(output.transpose(1, 2), input) + acc_train_loss += loss.item() * input.size(0) - 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) + nb_train_samples += input.size(0) - np = picoclvr.nb_properties(result_descr, height=self.height, width=self.width) + loss.backward() - acc_nb_requested_properties, _, acc_nb_missing_properties = zip(*np) - acc_nb_results = len(result_descr) + if nb_train_samples % args.batch_size == 0: + optimizer.step() - nb_requested_properties = sum(acc_nb_requested_properties) - nb_missing_properties = sum(acc_nb_missing_properties) + train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples)) - 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}%" - ) + log_string(f"train_perplexity {n_epoch} {train_perplexity}") - img = picoclvr.descr2img( - result_descr, [0], height=self.height, width=self.width - ) - 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}") +###################################################################### -###################################################################### +def run_tests(model, quizz_machine, deterministic_synthesis): + with torch.autograd.no_grad(): + model.eval() -log_string(f"device {device}") + nb_test_samples, acc_test_loss = 0, 0.0 + nb_samples_accumulated = 0 + for input in quizz_machine.batches(split="test"): + input = input.to(device) -def pruner_horizontal_green(p): - return not ("green" in p and ("left" in p or "right" in p)) + bs = model(mygpt.BracketedSequence(input)) + output = bs.x + loss = F.cross_entropy(output.transpose(1, 2), input) -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, -) + acc_test_loss += loss.item() * input.size(0) -vocabulary_size = task.vocabulary_size() + nb_test_samples += input.size(0) -log_string(f"vocabulary_size {vocabulary_size}") + main_test_accuracy = quizz_machine.produce_results( + n_epoch=n_epoch, + model=model, + result_dir=args.result_dir, + deterministic_synthesis=deterministic_synthesis, + ) -############################## + test_perplexity = math.exp(min(100, acc_test_loss / nb_test_samples)) -model = mygpt.MyGPT( - vocabulary_size=vocabulary_size, - dim_model=args.dim_model, - dim_keys=args.dim_keys, - dim_hidden=args.dim_hidden, - nb_heads=args.nb_heads, - nb_blocks=args.nb_blocks, - causal=True, - dropout=args.dropout, -) + log_string(f"test_perplexity {n_epoch} {test_perplexity}") -model.to(device) + model.main_test_accuracy = main_test_accuracy -nb_parameters = sum(p.numel() for p in model.parameters()) -log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)") ###################################################################### -nb_epochs_finished = 0 -if args.no_checkpoint: - log_string(f"not trying to load checkpoint.") +def create_c_quizzes( + models, + quizz_machine, + nb_for_train=1000, + nb_for_test=100, + min_ave_seq_logproba=None, +): + # We will store the generated quizzes for each number of + # correct prediction + recorded = dict([(n, []) for n in range(len(models) + 1)]) + + model_indexes = [] + sum_logits, sum_nb_c_quizzes = 0, 0 + + def nb_generated(): + return sum([sum([x.size(0) for x in recorded[n]]) for n in recorded.keys()]) + + def nb_validated(): + return sum( + [ + sum([x.size(0) for x in recorded[n]]) + for n in range(args.min_to_validate, args.max_to_validate + 1) + ] + ) -else: - try: - checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name) - checkpoint = torch.load(checkpoint_name) - nb_epochs_finished = checkpoint["nb_epochs_finished"] - model.load_state_dict(checkpoint["model_state"]) - torch.set_rng_state(checkpoint["rng_state"]) - if torch.cuda.is_available(): - torch.cuda.set_rng_state(checkpoint["cuda_rng_state"]) + nb_to_create = nb_for_train + nb_for_test + + while nb_validated() < nb_to_create: + ( + new_c_quizzes, + nb_correct, + ave_seq_logproba, + ) = quizz_machine.gang_create_c_quizzes( + nb=nb_to_create, + nb_models_for_generation=args.nb_models_for_generation, + models=models, + mode=args.generation_mode, + min_ave_seq_logproba=min_ave_seq_logproba, + n_epoch=n_epoch, + result_dir=args.result_dir, + ) - log_string(f"checkpoint loaded with {nb_epochs_finished} epochs finished.") + sum_logits += new_c_quizzes.size(0) * ave_seq_logproba + sum_nb_c_quizzes += new_c_quizzes.size(0) - except FileNotFoundError: - log_string("starting from scratch.") + if args.dirty_debug: + nb_correct = torch.randint( + len(models) + 1, nb_correct.size(), device=new_c_quizzes.device + ) - except: - log_string("error when loading the checkpoint.") - exit(1) + for n in range(nb_correct.max() + 1): + recorded[n].append(new_c_quizzes[nb_correct == n].clone()) -###################################################################### + nv = [recorded[n][-1].size(0) for n in recorded.keys()] -nb_epochs = args.nb_epochs if args.nb_epochs > 0 else nb_epochs_default + log_string(f"keep c_quizzes kept {nv} total {nb_validated()} / {nb_to_create}") -token_count = 0 -for input in task.batches(split="train"): - token_count += F.one_hot(input, num_classes=task.vocabulary_size()).sum((0, 1)) -token_probas = token_count / token_count.sum() -entropy = -torch.xlogy(token_probas, token_probas).sum() -train_set_perplexity = math.exp(entropy) + # concatenate and shuffle + for n in recorded.keys(): + if len(recorded[n]) > 0: + q = torch.cat(recorded[n], dim=0) + q = q[torch.randperm(q.size(0), device=q.device)] + recorded[n] = q + else: + del recorded[n] -############################## + new_c_quizzes = torch.cat( + [recorded[n] for n in range(args.min_to_validate, args.max_to_validate + 1)], + dim=0, + ) -if args.learning_rate_schedule == "cos": - learning_rate_schedule = {} - for n_epoch in range(args.nb_epochs): - u = n_epoch / args.nb_epochs * math.pi - learning_rate_schedule[n_epoch] = args.learning_rate * 0.5 * (1 + math.cos(u)) -else: - u = { - int(k): float(v) - for k, v in [ - tuple(x.split(":")) for x in args.learning_rate_schedule.split(",") + new_c_quizzes = new_c_quizzes[ + torch.randperm(new_c_quizzes.size(0), device=new_c_quizzes.device)[ + : nb_for_train + nb_for_test ] - } + ] - learning_rate_schedule = {} - learning_rate = args.learning_rate - for n_epoch in range(args.nb_epochs): - if n_epoch in u: - learning_rate = u[n_epoch] - learning_rate_schedule[n_epoch] = learning_rate + quizz_machine.store_c_quizzes(new_c_quizzes[:nb_for_train], for_train=True) + quizz_machine.store_c_quizzes(new_c_quizzes[nb_for_train:], for_train=False) -log_string(f"learning_rate_schedule {learning_rate_schedule}") + for n in recorded.keys(): + s = ( + "_validated" + if n >= args.min_to_validate and n <= args.max_to_validate + else "" + ) + quizz_machine.problem.save_quizzes( + recorded[n][:72], + args.result_dir, + f"culture_c_quiz_{n_epoch:04d}_N{n}{s}", + ) -############################## + return sum_logits / sum_nb_c_quizzes -nb_samples_seen = 0 -if nb_epochs_finished >= nb_epochs: - task.produce_results(nb_epochs_finished, model) +###################################################################### -for n_epoch in range(nb_epochs_finished, nb_epochs): +models = [] - learning_rate = learning_rate_schedule[n_epoch] +for k in range(args.nb_gpts): + model = mygpt.MyGPT( + vocabulary_size=vocabulary_size, + dim_model=args.dim_model, + dim_keys=args.dim_keys, + dim_hidden=args.dim_hidden, + nb_heads=args.nb_heads, + nb_blocks=args.nb_blocks, + causal=True, + dropout=args.dropout, + ).to(device) - log_string(f"learning_rate {learning_rate}") + model.main_test_accuracy = 0.0 + model.id = k - if args.optim == "sgd": - optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate) - elif args.optim == "adam": - optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) - elif args.optim == "adamw": - optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) - else: - raise ValueError(f"Unknown optimizer {args.optim}.") + models.append(model) - model.train() - nb_train_samples, acc_train_loss = 0, 0.0 +nb_parameters = sum(p.numel() for p in models[0].parameters()) +log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)") - 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) - 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() +min_ave_seq_logproba = None - with torch.autograd.no_grad(): +for n_epoch in range(args.nb_epochs): + log_string(f"--- epoch {n_epoch} ----------------------------------------") - model.eval() + a = [(model.id, float(model.main_test_accuracy)) for model in models] + a.sort(key=lambda p: p[0]) + s = " ".join([f"{p[1]*100:.02f}%" for p in a]) + log_string(f"current accuracies {s}") - nb_test_samples, acc_test_loss = 0, 0.0 + # select the model with lowest accuracy + models.sort(key=lambda model: model.main_test_accuracy) + model = models[0] - for input in task.batches(split="test"): - input = input.to(device) + log_string( + f"training model {model.id} main_test_accuracy {model.main_test_accuracy}" + ) - # input, loss_masks, true_images = task.excise_last_image(input) - # input, loss_masks = task.add_true_image(input, true_images, loss_masks) + # improve it + one_epoch(model, quizz_machine) - 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) + quizz_machine.renew_w_quizzes(args.nb_train_samples // args.nb_gpts) - 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"train_set_composition w_quizzes {quizz_machine.nb_batch_w_quizzes} c_quizzes {quizz_machine.nb_batch_c_quizzes}" + ) - log_string( - f"perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}" - ) + # test it + run_tests(model, quizz_machine, deterministic_synthesis=False) + + log_string( + f"test_set_composition w_quizzes {quizz_machine.nb_batch_w_quizzes} c_quizzes {quizz_machine.nb_batch_c_quizzes}" + ) - task.produce_results(n_epoch, model) + if min([m.main_test_accuracy for m in models]) >= args.accuracy_to_make_c_quizzes: + ave_seq_logproba = create_c_quizzes( + models, + quizz_machine, + nb_for_train=nb_new_c_quizzes_for_train, + nb_for_test=nb_new_c_quizzes_for_test, + min_ave_seq_logproba=min_ave_seq_logproba, + ) - checkpoint = { - "nb_epochs_finished": n_epoch + 1, - "model_state": model.state_dict(), - "rng_state": torch.get_rng_state(), - } + # We keep the first average logits as a reference + # if min_ave_seq_logproba is None: + # min_ave_seq_logproba = ave_seq_logproba + # else: + # log_string( + # f"min_ave_seq_logproba {min_ave_seq_logproba} ave_seq_logproba {ave_seq_logproba}" + # ) - if torch.cuda.is_available(): - checkpoint["cuda_rng_state"] = torch.cuda.get_rng_state() + # We update everyone + for model in models: + run_tests(model, quizz_machine, deterministic_synthesis=False) - checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name) - torch.save(checkpoint, checkpoint_name) - log_string(f"saved checkpoint {checkpoint_name}") ######################################################################