X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=main.py;h=6b00bbfd991178841468b75006f94121668c2b4f;hb=HEAD;hp=32447bf4b732f69c1d6b4d984e311805703f266c;hpb=87c9333e7800e62911cd4299500d4824d29a1ce1;p=culture.git diff --git a/main.py b/main.py index 32447bf..40772c2 100755 --- a/main.py +++ b/main.py @@ -3,191 +3,231 @@ # Any copyright is dedicated to the Public Domain. # https://creativecommons.org/publicdomain/zero/1.0/ -# Written by Francois Fleuret +# > A > f(A) > B ; > f(B) +# < f(B) ; < B < f(A) < A -# torch.backends.cuda.matmul.allow_tf23 -# torch.autocast(torch.bfloat16) +# Written by Francois Fleuret -import math, sys, argparse, time, tqdm, 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, grids, quiz_machine -if torch.cuda.is_available(): - device = torch.device("cuda") - torch.backends.cuda.matmul.allow_tf32 = True -else: - device = torch.device("cpu") +from quiz_machine import one_batch_masked_inplace_autoregression + +import threading, subprocess + +import torch.multiprocessing as mp ###################################################################### parser = argparse.ArgumentParser( - description="An implementation of GPT with cache.", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) -parser.add_argument( - "--task", - type=str, - default="picoclvr", - help="picoclvr, mnist, maze, snake, stack, expr", -) - -parser.add_argument("--log_filename", type=str, default="train.log", help=" ") +parser.add_argument("--log_filename", type=str, default="train.log") 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=None) +parser.add_argument("--resume", action="store_true", default=False) + +parser.add_argument("--max_percents_of_test_in_train", type=int, default=-1) + +parser.add_argument("--log_command", type=str, default=None) + +# ---------------------------------- + +parser.add_argument("--nb_epochs", type=int, default=10000) parser.add_argument("--batch_size", type=int, default=None) +parser.add_argument("--physical_batch_size", type=int, default=None) + +parser.add_argument("--inference_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("--optim", type=str, default="adam") +parser.add_argument("--nb_new_c_quizzes_for_train", type=int, default=None) + +parser.add_argument("--nb_new_c_quizzes_for_test", type=int, default=None) -parser.add_argument("--learning_rate", type=float, default=1e-4) +parser.add_argument("--learning_rate", type=float, default=5e-4) -parser.add_argument("--learning_rate_schedule", type=str, default="10: 2e-5,30: 4e-6") +parser.add_argument("--schedule_free", action="store_true", default=False) -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("--deterministic_synthesis", 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("--problem", type=str, default="grids") -parser.add_argument("--checkpoint_name", type=str, default="checkpoint.pth") +parser.add_argument("--nb_threads", type=int, default=1) -############################## -# picoclvr options +parser.add_argument("--gpus", type=str, default="all") -parser.add_argument("--picoclvr_nb_colors", type=int, default=5) +# ---------------------------------- -parser.add_argument("--picoclvr_height", type=int, default=12) +parser.add_argument("--nb_gpts", type=int, default=5) -parser.add_argument("--picoclvr_width", type=int, default=16) +parser.add_argument("--max_fail_to_validate", type=int, default=3) -parser.add_argument("--picocvlr_prune_properties", type=str, default="none") +parser.add_argument("--accuracy_to_make_c_quizzes", type=float, default=0.95) -############################## -# Maze options +parser.add_argument("--proba_understands", type=float, default=0.95) -parser.add_argument("--maze_height", type=int, default=13) +parser.add_argument("--proba_not_understands", type=float, default=0.1) -parser.add_argument("--maze_width", type=int, default=21) +parser.add_argument("--temperature_hot", type=float, default=1.5) -parser.add_argument("--maze_nb_walls", type=int, default=15) +parser.add_argument("--temperature_cold", type=float, default=1) -############################## -# Snake options +parser.add_argument("--prompt_noise", type=float, default=0.05) -parser.add_argument("--snake_height", type=int, default=6) +parser.add_argument("--dirty_debug", action="store_true", default=False) -parser.add_argument("--snake_width", type=int, default=8) +parser.add_argument("--test", type=str, default=None) -parser.add_argument("--snake_nb_colors", type=int, default=5) +###################################################################### -parser.add_argument("--snake_length", type=int, default=200) +grids_tasks = ", ".join( + [x.__name__.removeprefix("task_") for x in grids.Grids().all_tasks] +) -############################## -# Snake options +parser.add_argument( + "--grids_world_tasks", + type=str, + default="replace_color,translate,grow,frame", + help="A comma-separated subset of: " + grids_tasks + ".", +) -parser.add_argument("--stack_nb_steps", type=int, default=100) +parser.add_argument( + "--grids_science_tasks", + type=str, + default=None, + help="A comma-separated subset of: " + grids_tasks + ", or None.", +) -parser.add_argument("--stack_nb_stacks", type=int, default=1) +###################################################################### -parser.add_argument("--stack_nb_digits", type=int, default=3) +parser.add_argument("--sky_height", type=int, default=6) -parser.add_argument("--stack_fraction_values_for_train", type=float, default=None) +parser.add_argument("--sky_width", type=int, default=8) -############################## -# Expr options +parser.add_argument("--sky_nb_birds", type=int, default=3) -parser.add_argument("--expr_nb_variables", type=int, default=5) +parser.add_argument("--sky_nb_iterations", type=int, default=2) -parser.add_argument("--expr_sequence_length", type=int, default=30) +parser.add_argument("--sky_speed", type=int, default=3) ###################################################################### args = parser.parse_args() -assert args.picocvlr_prune_properties in {"none", "train+eval", "eval"} - if args.result_dir is None: - args.result_dir = f"results_{args.task}" + args.result_dir = f"results_culture" + +assert not args.grids_science_tasks or ( + len( + set(args.grids_world_tasks.split(",")) + & set(args.grids_science_tasks.split(",")) + ) + == 0 +), "World and science tasks have to be disjoint" ###################################################################### default_args = { - "picoclvr": { - "nb_epochs": 25, - "batch_size": 25, - "nb_train_samples": 250000, - "nb_test_samples": 10000, - }, - "mnist": { - "nb_epochs": 25, - "batch_size": 10, - "nb_train_samples": 250000, - "nb_test_samples": 10000, + "model": "37M", + "batch_size": 25, + "inference_batch_size": 50, + "nb_train_samples": 40000, + "nb_test_samples": 1000, +} + +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, }, - "maze": { - "nb_epochs": 25, - "batch_size": 25, - "nb_train_samples": 250000, - "nb_test_samples": 10000, + "4M": { + "dim_model": 256, + "dim_keys": 32, + "dim_hidden": 1024, + "nb_heads": 4, + "nb_blocks": 6, }, - "snake": { - "nb_epochs": 5, - "batch_size": 25, - "nb_train_samples": 250000, - "nb_test_samples": 10000, + "37M": { + "dim_model": 512, + "dim_keys": 64, + "dim_hidden": 2048, + "nb_heads": 8, + "nb_blocks": 12, }, - "stack": { - "nb_epochs": 5, - "batch_size": 25, - "nb_train_samples": 100000, - "nb_test_samples": 1000, + "122M": { + "dim_model": 768, + "dim_keys": 64, + "dim_hidden": 2048, + "nb_heads": 8, + "nb_blocks": 24, }, - "expr": { - "nb_epochs": 50, - "batch_size": 25, - "nb_train_samples": 250000, - "nb_test_samples": 10000, + "352M": { + "dim_model": 1024, + "dim_keys": 64, + "dim_hidden": 2048, + "nb_heads": 8, + "nb_blocks": 48, }, } -if args.task in default_args: - for k, v in default_args[args.task].items(): +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: +if args.resume: + assert os.path.isdir(args.result_dir) + +else: + try: + os.mkdir(args.result_dir) + except FileExistsError: print(f"result directory {args.result_dir} already exists") exit(1) @@ -215,1211 +255,1122 @@ def log_string(s): sys.stdout.flush() +###################################################################### +# Create a time-stamped archive of the source code + +with open("this_run.sh", "w") as f: + f.write(f"{' '.join(sys.argv)}\n") + +now = time.strftime("%Y%m%d-%H%M%S", time.localtime()) + +os.system(f"tar zcvf {args.result_dir}/src-{now}.tgz *.py *.sh") + +###################################################################### + +log_string(f"argv {' '.join(sys.argv)}") + for n in vars(args): log_string(f"args.{n} {getattr(args, n)}") + ###################################################################### +if args.gpus == "all": + gpus_idx = range(torch.cuda.device_count()) +else: + gpus_idx = [int(k) for k in args.gpus.split(",")] -# ra_mask is boolean, with 1s on the values to generate +gpus = [torch.device(f"cuda:{n}") for n in gpus_idx] +if torch.cuda.is_available(): + main_device = gpus[0] +else: + assert len(gpus) == 0 + main_device = torch.device("cpu") -def masked_inplace_autoregression( - model, - batch_size, - input, - ar_mask, - forbidden_tokens=None, - progress_bar_desc="autoregression", - device=torch.device("cpu"), -): - batches = zip(input.split(batch_size), ar_mask.split(batch_size)) +if args.dirty_debug: + args.nb_train_samples = 2500 + args.nb_test_samples = 100 - if progress_bar_desc is not None: - batches = tqdm.tqdm( - batches, - dynamic_ncols=True, - desc=progress_bar_desc, - total=input.size(0) // batch_size, +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, + max_nb_cached_chunks=len(gpus) * args.nb_train_samples // 100, + chunk_size=100, + nb_threads=args.nb_threads, + ) + +elif args.problem == "grids": + problem = grids.Grids( + max_nb_cached_chunks=len(gpus) * args.nb_train_samples // 100, + chunk_size=100, + nb_threads=args.nb_threads, + tasks=args.grids_world_tasks, + ) + + if args.grids_science_tasks is None: + science_w_quizzes = None + else: + science_problem = grids.Grids( + max_nb_cached_chunks=len(gpus) * args.nb_train_samples // 100, + chunk_size=100, + nb_threads=args.nb_threads, + tasks=args.grids_science_tasks, ) + science_w_quizzes = science_problem.generate_w_quizzes(100) - for input, ar_mask in batches: - 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 not args.resume: + science_problem.save_some_examples(args.result_dir, "science_") +else: + raise ValueError + +if not args.resume: + problem.save_some_examples(args.result_dir) + +quiz_machine = quiz_machine.QuizMachine( + problem=problem, + batch_size=args.inference_batch_size, + result_dir=args.result_dir, + prompt_noise=args.prompt_noise, + logger=log_string, + device=main_device, +) + ###################################################################### +log_string(f"main_device {main_device} gpus {[ str(g) for g in gpus]}") -class Task: - def batches(self, split="train"): - pass +vocabulary_size = quiz_machine.vocabulary_size() - def vocabulary_size(self): - pass +log_string(f"vocabulary_size {vocabulary_size}") + +###################################################################### - def produce_results(self, n_epoch, model): - pass + +def optimizer_to(optim, device): + for param in optim.state.values(): + # Not sure there are any global tensors in the state dict + if isinstance(param, torch.Tensor): + param.data = param.data.to(device) + if param._grad is not None: + param._grad.data = param._grad.data.to(device) + elif isinstance(param, dict): + for subparam in param.values(): + if isinstance(subparam, torch.Tensor): + subparam.data = subparam.data.to(device) + if subparam._grad is not None: + subparam._grad.data = subparam._grad.data.to(device) ###################################################################### -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, - progress_bar_desc=None, - device=self.device, - ) - model.train(t) - input, loss_masks = self.trim((input, loss_masks)) +def run_tests(model, quiz_machine, local_device=main_device): + with torch.autograd.no_grad(): + model.to(local_device).eval() + if args.schedule_free: + model.optimizer.eval() - return input, loss_masks + nb_test_samples, acc_test_loss = 0, 0.0 + nb_samples_accumulated = 0 - ###################### + full_input, full_mask_loss = quiz_machine.data_input(model, split="test") + src = zip( + full_input.split(args.batch_size), full_mask_loss.split(args.batch_size) + ) - def __init__( - self, - nb_train_samples, - nb_test_samples, - batch_size, - height, - width, - nb_colors=5, - 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, + for input, mask_loss in tqdm.tqdm( + src, + dynamic_ncols=True, + desc="test", + total=full_input.size(0) // args.batch_size, + ): + input = input.to(local_device) + mask_loss = mask_loss.to(local_device) + targets = input + + output = model(mygpt.BracketedSequence(input)).x + loss_per_token = F.cross_entropy( + output.transpose(1, 2), targets, reduction="none" ) + loss = (loss_per_token * mask_loss).mean() + acc_test_loss += loss.item() * input.size(0) + nb_test_samples += input.size(0) + + test_perplexity = math.exp(min(100, acc_test_loss / nb_test_samples)) - self.height = height - self.width = width - self.batch_size = batch_size - self.device = device - self.pruner_train = pruner_train - self.pruner_eval = pruner_eval - - param = { - "nb_train_samples": nb_train_samples, - "nb_test_samples": nb_test_samples, - "height": height, - "width": width, - "nb_colors": nb_colors, - "batch_size": batch_size, - "rng_state": list(torch.get_rng_state()), - } + log_string(f"test_perplexity {n_epoch} model {model.id} {test_perplexity}") - log_string( - f"generating {nb_train_samples+nb_test_samples} samples (can take some time)" - ) - self.train_descr = generate_descr( - nb_train_samples, "train", pruner=self.pruner_train + model.main_test_accuracy = quiz_machine.produce_results( + n_epoch=n_epoch, + model=model, + input=full_input[:2000], + result_dir=args.result_dir, ) - self.test_descr = generate_descr(nb_test_samples, "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}" - ): - yield self.trim(batch) - def vocabulary_size(self): - 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 +###################################################################### - 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) +def one_epoch(model, quiz_machine, local_device=main_device): + model.to(local_device).train() + optimizer_to(model.optimizer, local_device) - 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}%" + if args.schedule_free: + model.optimizer.train() + + nb_train_samples, acc_train_loss = 0, 0.0 + + hard_w_quizzes = [] + + full_input, full_mask_loss = quiz_machine.data_input(model, split="train") + src = zip(full_input.split(args.batch_size), full_mask_loss.split(args.batch_size)) + + for input, mask_loss in tqdm.tqdm( + src, + dynamic_ncols=True, + desc="training", + total=full_input.size(0) // args.batch_size, + ): + input = input.to(local_device) + mask_loss = mask_loss.to(local_device) + + if nb_train_samples % args.batch_size == 0: + model.optimizer.zero_grad() + + targets = input + + output = model(mygpt.BracketedSequence(input)).x + loss_per_token = F.cross_entropy( + output.transpose(1, 2), targets, reduction="none" ) + loss = (loss_per_token * mask_loss).mean() + model.loss + acc_train_loss += loss.item() * input.size(0) - ###################################################################### + loss_per_samples = loss_per_token.detach().flatten(1).mean(dim=1) - def produce_results(self, n_epoch, model): - self.compute_missing_properties(n_epoch, model) + nb_train_samples += input.size(0) - if self.pruner_eval is not None: - self.compute_missing_properties(n_epoch, model, self.pruner_eval) + loss.backward() - nb_tokens_to_generate = self.height * self.width + 3 - result_descr = [] - nb_per_primer = 8 - primer = [] + if nb_train_samples % args.batch_size == 0: + model.optimizer.step() - 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 + train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples)) - 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) + log_string(f"train_perplexity {n_epoch} model {model.id} {train_perplexity}") - np = picoclvr.nb_properties(result_descr, height=self.height, width=self.width) + run_tests(model, quiz_machine) - acc_nb_requested_properties, _, acc_nb_missing_properties = zip(*np) - acc_nb_results = len(result_descr) + # threshold = torch.cat([l for _, l in hard_w_quizzes], dim=0).sort().values + # threshold = threshold[threshold.size(0) // 2] - nb_requested_properties = sum(acc_nb_requested_properties) - nb_missing_properties = sum(acc_nb_missing_properties) + # model.hard_w_quizzes = torch.cat( + # [x[l >= threshold] for x, l in hard_w_quizzes], dim=0 + # ) - 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}%" - ) + model.to(main_device) + optimizer_to(model.optimizer, main_device) - img = picoclvr.descr2img(result_descr, 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, - ) +###################################################################### + + +def model_transformer_hot(model): + model.temperature = args.temperature_hot + # model.set_noise_injection(1.0, ("ffw", args.nb_blocks // 2)) - 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=0.0 - ) - log_string(f"wrote {image_name}") +def model_transformer_cold(model): + model.temperature = args.temperature_cold + # pass + + +c_quizzes_procedure = [ + (("f_B", "f_A", "A", "B"), (1, 0, 0, 0), model_transformer_hot), + (("f_B", "f_A", "A", "B"), (0, 1, 1, 1), model_transformer_cold), + (("A", "f_A", "B", "f_B"), (0, 0, 0, 1), model_transformer_cold), + (("f_A", "A", "f_B", "B"), (0, 0, 0, 1), model_transformer_cold), +] ###################################################################### -class TaskMNIST(Task): - def __init__(self, batch_size, device=torch.device("cpu")): - self.device = device - self.batch_size = batch_size +def save_additional_results(model, models, science_w_quizzes): + # Save generated quizzes with the successive steps + + recorder = [] + + c_quizzes = quiz_machine.generate_c_quizzes( + 64, + model_for_generation=model, + procedure=c_quizzes_procedure, + recorder=recorder, + ) + + # This is nb_quizzes x nb_models + + seq_logproba = quiz_machine.models_logprobas( + models, c_quizzes, ("A", "f_A", "B", "f_B"), (0, 0, 0, 1), (0, 0, 1, 0) + ) + quiz_machine.models_logprobas( + models, c_quizzes, ("f_A", "A", "f_B", "B"), (0, 0, 0, 1), (0, 0, 1, 0) + ) + + probas = seq_logproba.exp() + + comments = [] + + for l in seq_logproba: + comments.append("proba " + " ".join([f"{x.exp().item():.02f}" for x in l])) - def batches(self, split="train"): - assert split in {"train", "test"} - data_set = torchvision.datasets.MNIST( - root="./data", train=(split == "train"), download=True + ## + + c_quizzes = torch.cat([c[:, None, :] for c, _, in recorder], dim=1) + predicted_parts = torch.cat([t[:, None, :] for _, t in recorder], dim=1) + nb_steps = c_quizzes.size(1) + c_quizzes = c_quizzes.reshape(-1, c_quizzes.size(-1)) + predicted_parts = predicted_parts.reshape(-1, predicted_parts.size(-1)) + + # We have comments only for the final quiz, not the successive + # steps, so we have to add nb_steps-1 empty comments + + steps_comments = [] + for c in comments: + steps_comments += [""] * (nb_steps - 1) + [c] + + filename = f"non_validated_{n_epoch:04d}_{model.id:02d}.png" + + quiz_machine.problem.save_quizzes_as_image( + args.result_dir, + filename, + quizzes=c_quizzes, + predicted_parts=predicted_parts, + comments=steps_comments, + nrow=nb_steps * 2, # two quiz per row + ) + + log_string(f"wrote {filename}") + + ###################################################################### + + if science_w_quizzes is not None: + struct = ("A", "f_A", "B", "f_B") + mask = (0, 0, 0, 1) + result, correct = quiz_machine.predict( + model=model, + quizzes=science_w_quizzes.to(main_device), + struct=struct, + mask=mask, ) - 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 + predicted_parts = torch.tensor(mask, device=correct.device)[None, :].expand( + correct.size(0), -1 + ) + correct = (2 * correct - 1) * (predicted_parts.sum(dim=-1) == 1).long() - 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 + nb_correct = (correct == 1).long().sum() + nb_total = (correct != 0).long().sum() + + log_string( + f"science_accuracy {n_epoch} model {model.id} val {nb_correct} / {nb_total}" ) - 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, + + i = correct == 1 + j = correct != 1 + + result = torch.cat([result[i], result[j]], dim=0) + correct = torch.cat([correct[i], correct[j]], dim=0) + correct_parts = predicted_parts * correct[:, None] + + result = result[:128] + predicted_parts = predicted_parts[:128] + correct_parts = correct_parts[:128] + + quiz_machine.problem.save_quizzes_as_image( + args.result_dir, + f"culture_science_{n_epoch:04d}_{model.id:02d}.png", + quizzes=result, + predicted_parts=predicted_parts, + correct_parts=correct_parts, ) - log_string(f"wrote {image_name}") ###################################################################### -import maze +def record_new_c_quizzes(models, quiz_machine, nb_for_train=1000, nb_for_test=100): + nb_to_validate = nb_for_train + nb_for_test + nb_to_generate_per_iteration = max(args.physical_batch_size, nb_to_validate) + nb_validated = 0 -class TaskMaze(Task): - def map2seq(self, *m): - return torch.cat([x.flatten(1) for x in m], 1) + recorded_validated = [] - 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))) + start_time = time.perf_counter() - 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"), + nb_validated_per_model = torch.zeros(len(models), dtype=torch.int64) + + while nb_validated_per_model.sum() < nb_to_validate: + # We use the model that has generated the fewest quizzes to + # balance the number of quizzes per model overall + + # model_for_generation = sorted( + # models, key=lambda m: nb_validated_per_model[m.id] + # )[0] + + model_for_generation = models[torch.randint(len(models), (1,)).item()] + + # We generate quizzes with a procedure that injects some + # structured noise + + c_quizzes = quiz_machine.generate_c_quizzes( + nb_to_generate_per_iteration, + model_for_generation=model, + procedure=c_quizzes_procedure, ) - 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"), + + # We discard the trivial ones, according to a criterion + # specific to the world quizzes (e.g. B=f(B)) + + to_keep = quiz_machine.problem.trivial(c_quizzes) == False + + c_quizzes = c_quizzes[to_keep] + + # This is nb_quizzes x nb_models + + seq_logproba = quiz_machine.models_logprobas( + models, c_quizzes, ("A", "f_A", "B", "f_B"), (0, 0, 0, 1), (0, 0, 1, 0) + ) + quiz_machine.models_logprobas( + models, c_quizzes, ("f_A", "A", "f_B", "B"), (0, 0, 0, 1), (0, 0, 1, 0) ) - 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, + + probas = seq_logproba.exp() + + nb_succeed = (probas >= args.proba_understands).long().sum(dim=1) + nb_fail = (probas <= args.proba_not_understands).long().sum(dim=1) + + to_keep = ( + (nb_succeed + nb_fail == probas.size(1)) + & (nb_fail >= 1) + & (nb_fail <= args.max_fail_to_validate) ) - 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 + c_quizzes = c_quizzes[to_keep] - if count.max() == 0: - count = None + if c_quizzes.size(0) > 0: + nb_validated_per_model[model_for_generation.id] += c_quizzes.size(0) + recorded_validated.append(c_quizzes) + nb_validated = c_quizzes.size(0) else: - count = count[ - : count.sum(1).nonzero().max() + 1, : count.sum(0).nonzero().max() + 1 - ] + nb_validated = 0 - return nb_total, nb_correct, count + total_nb_validated = nb_validated_per_model.sum().item() - def produce_results(self, n_epoch, model): - with torch.autograd.no_grad(): - t = model.training - model.eval() + duration = time.perf_counter() - start_time - train_nb_total, train_nb_correct, count = self.compute_error( - model, "train", nb_to_use=1000 - ) - log_string( - f"accuracy_train {n_epoch} nb_total {train_nb_total} nb_correct {train_nb_correct} accuracy {(100.0*train_nb_correct)/train_nb_total:.02f}%" - ) + if total_nb_validated > 0: + if total_nb_validated < nb_to_validate: + d = ( + (nb_to_validate - total_nb_validated) + * duration + / total_nb_validated + ) + e = (datetime.datetime.now() + datetime.timedelta(seconds=d)).strftime( + "%a %H:%M" + ) + else: + e = "now!" + else: + e = "???" - test_nb_total, test_nb_correct, count = self.compute_error( - model, "test", nb_to_use=1000 - ) - log_string( - f"accuracy_test {n_epoch} nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%" - ) + log_string( + f"keep c_quizzes model {model_for_generation.id} validated {nb_validated} / {nb_to_generate_per_iteration} ({100*nb_validated/nb_to_generate_per_iteration:.02f}%) nb_accumulated {total_nb_validated} / {nb_to_validate} (finishes {e} -- {int((total_nb_validated * 3600)/duration)}/h)" + ) - 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 - ) + validated_quizzes = torch.cat(recorded_validated, dim=0) - 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}") + ###################################################################### + # store the new c_quizzes which have been validated - model.train(t) + v_train = validated_quizzes[:nb_for_train] + quiz_machine.store_c_quizzes(v_train, for_train=True) + v_test = validated_quizzes[nb_for_train:nb_to_validate] + quiz_machine.store_c_quizzes(v_test, for_train=False) -###################################################################### + ###################################################################### + # save images + vq = validated_quizzes[torch.randperm(validated_quizzes.size(0))[:128]] -import snake + if vq.size(0) > 0: + seq_logproba = quiz_machine.models_logprobas( + models, vq, ("A", "f_A", "B", "f_B"), (0, 0, 0, 1), (0, 0, 1, 0) + ) + quiz_machine.models_logprobas( + models, vq, ("f_A", "A", "f_B", "B"), (0, 0, 0, 1), (0, 0, 1, 0) + ) + probas = seq_logproba.exp() -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, + comments = [] + + for l in seq_logproba: + comments.append("proba " + " ".join([f"{x.exp().item():.02f}" for x in l])) + + filename = f"culture_c_quiz_{n_epoch:04d}.png" + quiz_machine.problem.save_quizzes_as_image( + args.result_dir, filename, vq, comments=comments ) - self.test_input, self.test_prior_visits, _, _ = snake.generate_sequences( - nb_test_samples, - height, - width, - nb_colors, - length, - prompt_length, - self.device, + + +###################################################################### + +# The generator is very similar to a "solving GPT" except that it +# deals with quizzes prologued with one token per solving GPT that +# indicates if the said model solves it or not. +# +# There are three levels of solving 0->proba<=proba_not_understands, +# 2->proba>=proba_understands and 1 otherwise. + + +def generate_c_quizzes_with_generator(generator, quiz_machine, nb): + generator.to(main_device) + + struct = ("A", "f_A", "B", "f_B") + + c_quizzes = quiz_machine.problem.create_empty_quizzes(nb, struct=struct) + ar_mask = quiz_machine.make_quiz_mask(c_quizzes, struct, (1, 1, 1, 1)) + + i = F.one_hot( + torch.randint(args.nb_gpts, (c_quizzes.size(0),)), + num_classes=args.nb_gpts, + ) + + prologs_c_quizzes = token_prolog_0 * i + token_prolog_2 * (1 - i) + prologs_ar_mask = ar_mask.new_zeros(ar_mask.size(0), prologs_c_quizzes.size(1)) + + prologued_c_quizzes = torch.cat([prologs_c_quizzes, c_quizzes], dim=1).to( + main_device + ) + prologued_ar_mask = torch.cat([prologs_ar_mask, ar_mask], dim=1).to(main_device) + + seq_logproba = torch.zeros( + prologued_c_quizzes.size(0), device=prologued_c_quizzes.device + ) + + generator.temperature = args.temperature_hot + + with torch.autograd.no_grad(): + t = generator.training + generator.eval() + + one_batch_masked_inplace_autoregression( + generator, + prologued_c_quizzes, + prologued_ar_mask, + seq_logproba, + deterministic_synthesis=False, ) - 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 + generator.train(t) + + generator.reset_transformations() - # snake.solver(result,ar_mask) + prologued_c_quizzes = ( + prologued_c_quizzes * (prologued_c_quizzes < vocabulary_size).long() + ) - masked_inplace_autoregression( - model, self.batch_size, result, ar_mask, device=self.device + c_quizzes = prologued_c_quizzes[:, prologs_c_quizzes.size(1) :] + + return c_quizzes.to("cpu"), prologs_c_quizzes.to("cpu") + + +def batches_for_generator(generator, quiz_machine, models, fraction_w_quizzes=1.0): + samples = [] + + for _ in range(args.nb_train_samples // args.batch_size): + while sum([x.size(0) for x in samples]) < args.batch_size: + # Generate a bunch of quizzes + + if torch.rand(1).item() <= fraction_w_quizzes: + # Either we start with the world quizzes + c_quizzes = quiz_machine.problem.generate_w_quizzes( + args.batch_size, progress_bar=False + ) + else: + # Or we use the generator itself to generate them + c_quizzes, _ = generate_c_quizzes_with_generator( + generator, quiz_machine, args.batch_size ) - nb_total = ((prior_visits > 0) * ar_mask).sum() + # We remove the trivial ones + to_keep = quiz_machine.problem.trivial(c_quizzes) == False + c_quizzes = c_quizzes[to_keep] + + # If there are remaining ones, we compute the true prolog + # that indicates how the GPTs solve it + + if c_quizzes.size(0) > 0: + seq_logproba = quiz_machine.models_logprobas( + models, + c_quizzes, + ("A", "f_A", "B", "f_B"), + (0, 0, 0, 1), + (0, 0, 1, 0), + ) + quiz_machine.models_logprobas( + models, + c_quizzes, + ("f_A", "A", "f_B", "B"), + (0, 0, 0, 1), + (0, 0, 1, 0), + ) - nb_correct = ( - (result == input).long() * (prior_visits > 0) * ar_mask - ).sum() + probas = seq_logproba.exp() - # nb_total = result.size(0) - # nb_correct = ((result - input).abs().sum(1) == 0).sum() + u0 = probas <= args.proba_not_understands + u2 = probas >= args.proba_understands + u1 = (u0 | u2) == False - return nb_total, nb_correct + prologs = ( + (u0.long() * token_prolog_0) + + (u1.long() * token_prolog_1) + + (u2.long() * token_prolog_2) + ) - # train_nb_total, train_nb_correct = compute_nb_correct( - # self.train_input, self.train_prior_visits - # ) + prologued_c_quizzes = torch.cat([prologs, c_quizzes], dim=1) - # 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}%" - # ) + # nb_u2 = u2.long().sum(dim=1) + # nb_u0 = u0.long().sum(dim=1) + # prologued_c_quizzes = prologued_c_quizzes[(nb_u2 >= 1) & (nb_u0 >= 1)] - test_nb_total, test_nb_correct = compute_nb_correct( - self.test_input[:1000], self.test_prior_visits[:1000] - ) + if prologued_c_quizzes.size(0) > 0: + samples.append(prologued_c_quizzes) - log_string( - f"accuracy_test {n_epoch} nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%" - ) + # Now we yield a batch - model.train(t) + x = torch.cat(samples, dim=0) + samples = [x[args.batch_size :]] + yield x[: args.batch_size] -###################################################################### +def one_generator_epoch( + generator, quiz_machine, models, fraction_w_quizzes, local_device=main_device +): + model.to(local_device).train() + + optimizer = torch.optim.Adam(generator.parameters(), lr=args.learning_rate) -import stack + nb_train_samples, acc_train_loss = 0, 0.0 + src = batches_for_generator( + generator=generator, + quiz_machine=quiz_machine, + models=models, + fraction_w_quizzes=fraction_w_quizzes, + ) -class TaskStack(Task): - def __init__( - self, - nb_train_samples, - nb_test_samples, - batch_size, - nb_steps, - nb_stacks, - nb_digits, - fraction_values_for_train=None, - device=torch.device("cpu"), + for input in tqdm.tqdm( + src, + dynamic_ncols=True, + desc="training", + total=args.nb_train_samples // args.batch_size, ): - self.batch_size = batch_size - self.nb_steps = nb_steps - self.nb_stacks = nb_stacks - self.nb_digits = nb_digits - self.device = device - - if fraction_values_for_train is None: - values_for_train = None - values_for_test = None - else: - all = torch.randperm(10**nb_digits) - nb_for_train = int(all.size(0) * fraction_values_for_train) - values_for_train = all[:nb_for_train] - values_for_test = all[nb_for_train:] - - self.train_input, self.train_stack_counts = stack.generate_sequences( - nb_train_samples, - nb_steps, - nb_stacks, - nb_digits, - values_for_train, - self.device, - ) + input = input.to(local_device) - self.test_input, self.test_stack_counts = stack.generate_sequences( - nb_test_samples, - nb_steps, - nb_stacks, - nb_digits, - values_for_test, - self.device, - ) + if nb_train_samples % args.batch_size == 0: + optimizer.zero_grad() - i = torch.logical_and(self.test_input % 2 == 1, self.test_input < 2 * nb_stacks) - counts = self.test_stack_counts.flatten()[i.flatten()] - counts = F.one_hot(counts).sum(0) - log_string(f"test_pop_stack_counts {counts}") - - 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): - result = input.clone() - stack.remove_popped_values(result, self.nb_stacks, self.nb_digits) - ar_mask = (result != input).long() - masked_inplace_autoregression( - model, self.batch_size, result, ar_mask, device=self.device - ) + targets = input - errors = ((result != input).long() * ar_mask).reshape( - -1, 1 + self.nb_digits - ) - ar_mask = ar_mask.reshape(-1, 1 + self.nb_digits) + output = generator(mygpt.BracketedSequence(input)).x + loss = F.cross_entropy(output.transpose(1, 2), targets) + acc_train_loss += loss.item() * input.size(0) + nb_train_samples += input.size(0) - nb_total = ar_mask.max(1).values.sum() - nb_correct = nb_total - errors.max(1).values.sum() + loss.backward() - return nb_total, nb_correct + if nb_train_samples % args.batch_size == 0: + optimizer.step() - test_nb_total, test_nb_correct = compute_nb_correct(self.test_input[:1000]) + train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples)) - log_string( - f"accuracy_test {n_epoch} nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%" - ) + log_string(f"train_perplexity {n_epoch} generator - {train_perplexity}") - ############################################################## - # Log a few generated sequences - input = self.test_input[:10, : 12 * (1 + self.nb_digits)] - result = input.clone() - stack.remove_popped_values(result, self.nb_stacks, self.nb_digits) - ar_mask = (result != input).long() - for n in range(result.size(0)): - log_string( - f"test_before {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}" + generator.to(main_device) + + +###################################################################### + + +def train_complexifier(model_gen, model_pred1, model_pred2): + samples = [] + perf = [] + + optimizer = torch.optim.Adam(model_gen.parameters(), lr=args.learning_rate) + + nb_train_samples, acc_train_loss = 0, 0.0 + + for n_epoch in range(args.nb_epochs): + for b in range(args.nb_train_samples // args.batch_size): + while sum([x.size(0) for x in samples]) < args.batch_size: + c_quizzes = quiz_machine.generate_c_quizzes( + args.inference_batch_size, + model_for_generation=model_gen, + procedure=c_quizzes_procedure, ) - masked_inplace_autoregression( - model, self.batch_size, result, ar_mask, device=self.device + to_keep = quiz_machine.problem.trivial(c_quizzes) == False + c_quizzes = c_quizzes[to_keep] + if c_quizzes.size(0) > 0: + seq_logproba = quiz_machine.models_logprobas( + [model_pred1, model_pred2], + c_quizzes, + ("A", "f_A", "B", "f_B"), + (0, 0, 0, 1), + ) + quiz_machine.models_logprobas( + [model_pred1, model_pred2], + c_quizzes, + ("f_A", "A", "f_B", "B"), + (0, 0, 0, 1), + ) + probas = seq_logproba.exp() + to_keep = (probas[:, model_pred1.id] >= args.proba_understands) & ( + probas[:, model_pred2.id] <= args.proba_not_understands + ) + log_string( + f"generating {to_keep.long().sum()} / {c_quizzes.size(0)}" + ) + c_quizzes = c_quizzes[to_keep] + if c_quizzes.size(0): + samples.append(c_quizzes) + + log_string(f"full batch {sum([x.size(0) for x in samples])}") + + x = torch.cat(samples, dim=0) + + input = x[: args.batch_size] + samples = [x[args.batch_size :]] + + # ------------------- + + seq_logproba = quiz_machine.models_logprobas( + [model_pred1, model_pred2], + input, + ("A", "f_A", "B", "f_B"), + (0, 0, 0, 1), + ) + quiz_machine.models_logprobas( + [model_pred1, model_pred2], + input, + ("f_A", "A", "f_B", "B"), + (0, 0, 0, 1), ) - for n in range(result.size(0)): - log_string( - f"test_after {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}" + + comments = [] + + for l in seq_logproba: + comments.append( + f"proba {l[model_pred1.id].exp().item():.02f} {l[model_pred2.id].exp().item():.02f}" ) - ############################################################## - model.train(t) + filename = f"batch_{n_epoch:04d}_{b:04d}.png" + quiz_machine.problem.save_quizzes_as_image( + args.result_dir, filename, input, comments=comments + ) + log_string(f"wrote {filename}") + + # ------------------------ + + input = input.to(main_device) + + if nb_train_samples % args.batch_size == 0: + optimizer.zero_grad() + + output = model_gen(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) + + loss.backward() + + if nb_train_samples % args.batch_size == 0: + optimizer.step() + + train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples)) + + log_string(f"train_perplexity {n_epoch} model ae {train_perplexity}") ###################################################################### +models = [] -import expr +def compute_causal_attzero(t_q, t_k): + return t_q < t_k -class TaskExpr(Task): - def __init__( - self, - nb_train_samples, - nb_test_samples, - nb_variables, - sequence_length, - batch_size, - device=torch.device("cpu"), - ): - self.batch_size = batch_size - self.device = device - - train_sequences = expr.generate_sequences( - nb_train_samples, - nb_variables=nb_variables, - length=sequence_length, - # length=2 * sequence_length, - # randomize_length=True, - ) - test_sequences = expr.generate_sequences( - nb_test_samples, - nb_variables=nb_variables, - length=sequence_length, - ) - self.char2id = dict( - [ - (c, n) - for n, c in enumerate( - set("#" + "".join(train_sequences + test_sequences)) - ) - ] + +if args.schedule_free: + import schedulefree + +for k in range(args.nb_gpts): + log_string(f"creating model {k} and its w_quizzes") + + 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, + compute_attzero=compute_causal_attzero, + dropout=args.dropout, + ).to(main_device) + + model.id = k + + if args.schedule_free: + model.optimizer = schedulefree.AdamWScheduleFree( + model.parameters(), lr=args.learning_rate ) - self.id2char = dict([(n, c) for c, n in self.char2id.items()]) - - self.filler, self.space = self.char2id["#"], self.char2id[" "] - - len_max = max([len(x) for x in train_sequences]) - self.train_input = torch.cat( - [ - torch.tensor( - [ - [self.char2id[c] for c in s + "#" * (len_max - len(s))] - for s in train_sequences - ] - ) - ], - 0, - ).to(device) - - len_max = max([len(x) for x in test_sequences]) - self.test_input = torch.cat( - [ - torch.tensor( - [ - [self.char2id[c] for c in s + "#" * (len_max - len(s))] - for s in test_sequences - ] - ) - ], - 0, - ).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 - ): - if split == "train": - last = (batch != self.filler).max(0).values.nonzero().max() + 1 - batch = batch[:, :last] - yield batch - - def vocabulary_size(self): - return self.nb_codes - - def seq2str(self, s): - return "".join([self.id2char[k.item()] for k in s]) - - def produce_results(self, n_epoch, model): - with torch.autograd.no_grad(): - t = model.training - model.eval() - - def compute_nb_correct(input): - result = input.clone() - ar_mask = (result == self.space).long().cumsum(dim=1).clamp(max=1) - result = (1 - ar_mask) * result + ar_mask * self.filler - masked_inplace_autoregression( - model, self.batch_size, result, ar_mask, device=self.device - ) + else: + model.optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate) - nb_total = input.size(0) - nb_correct = (input == result).long().min(1).values.sum() + model.main_test_accuracy = 0.0 - ####################################################################### - # Comput predicted vs. true variable values + model.train_w_quizzes = quiz_machine.problem.generate_w_quizzes( + args.nb_train_samples + ) - nb_delta = torch.zeros(5, dtype=torch.int64) - nb_missed = 0 + model.test_w_quizzes = quiz_machine.problem.generate_w_quizzes(args.nb_test_samples) - values_input = expr.extract_results([self.seq2str(s) for s in input]) - values_result = expr.extract_results([self.seq2str(s) for s in result]) + models.append(model) - for i, r in zip(values_input, values_result): - for n, vi in i.items(): - vr = r.get(n) - if vr is None or vr < 0: - nb_missed += 1 - else: - d = abs(vr-vi) - if d >= nb_delta.size(0): - nb_missed += 1 - else: - nb_delta[d] += 1 +###################################################################### - ###################################################################### +if args.test == "quant": + nb_bits = 8 + for model in models: + model.trunk.insert( + 12, + mygpt.CacheWrapper( + mygpt.RandomBypass( + nn.Sequential( + nn.Linear(args.dim_model, nb_bits), + mygpt.BSQ(nb_bits), + nn.Linear(nb_bits, args.dim_model), + ), + 0.1, + ) + ), + ) - return nb_total, nb_correct, nb_delta, nb_missed + print(model) + exit(0) - test_nb_total, test_nb_correct, test_nb_delta, test_nb_missed = compute_nb_correct(self.test_input[:1000]) - log_string( - f"accuracy_test {n_epoch} nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%" - ) +###################################################################### - nb_total = test_nb_delta.sum() + test_nb_missed - for d in range(test_nb_delta.size(0)): - log_string(f"error_value {n_epoch} delta {d} {test_nb_delta[d]} {test_nb_delta[d]*100/nb_total:.02f}%") - log_string(f"error_value {n_epoch} missed {test_nb_missed} {test_nb_missed*100/nb_total:.02f}%") - - - ############################################################## - # Log a few generated sequences - input = self.test_input[:10] - result = input.clone() - ar_mask = (result == self.space).long().cumsum(dim=1).clamp(max=1) - result = (1 - ar_mask) * result + ar_mask * self.filler - for n in range(result.size(0)): - log_string(f"test_before {self.seq2str(result[n])}") - masked_inplace_autoregression( - model, self.batch_size, result, ar_mask, device=self.device - ) - correct = (1 - ar_mask) * self.space + ar_mask * input - for n in range(result.size(0)): - comment = "GOOD" if (result[n] - input[n]).abs().max() == 0 else "" - log_string(f"test_after {self.seq2str(result[n])} {comment}") - log_string(f"correct {self.seq2str(correct[n])}") - ############################################################## +current_epoch = 0 + +if args.resume: + for model in models: + filename = f"gpt_{model.id:03d}.pth" - model.train(t) + try: + d = torch.load(os.path.join(args.result_dir, filename)) + model.load_state_dict(d["state_dict"]) + model.optimizer.load_state_dict(d["optimizer_state_dict"]) + model.main_test_accuracy = d["main_test_accuracy"] + log_string(f"successfully loaded {filename}") + except FileNotFoundError: + log_string(f"cannot find {filename}") + pass + try: + filename = "c_quizzes.pth" + quiz_machine.load_c_quizzes(os.path.join(args.result_dir, filename)) + log_string(f"successfully loaded {filename}") + except FileNotFoundError: + log_string(f"cannot find {filename}") + pass + + try: + filename = "state.pth" + state = torch.load(os.path.join(args.result_dir, filename)) + log_string(f"successfully loaded {filename}") + current_epoch = state["current_epoch"] + except FileNotFoundError: + log_string(f"cannot find {filename}") + pass ###################################################################### +nb_parameters = sum(p.numel() for p in models[0].parameters()) +log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)") -def picoclvr_pruner_horizontal_green(p): - return not ("green" in p and ("left" in p or "right" in p)) +###################################################################### +if args.nb_new_c_quizzes_for_train is None: + args.nb_new_c_quizzes_for_train = args.nb_train_samples // 100 -picoclvr_pruner_train = ( - picoclvr_pruner_horizontal_green - if args.picocvlr_prune_properties in {"train+eval"} - else None -) +if args.nb_new_c_quizzes_for_test is None: + args.nb_new_c_quizzes_for_test = args.nb_test_samples // 100 -picoclvr_pruner_eval = ( - (lambda p: not picoclvr_pruner_horizontal_green(p)) - if args.picocvlr_prune_properties in {"train+eval", "eval"} - else None +log_string( + f"nb_new_c_quizzes_for_train {args.nb_new_c_quizzes_for_train} nb_new_c_quizzes_for_test {args.nb_new_c_quizzes_for_test}" ) ###################################################################### -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, - ) +if args.dirty_debug: + args.accuracy_to_make_c_quizzes = 0.0 + args.nb_gpts = 2 + args.nb_new_c_quizzes_for_train = 100 + args.nb_new_c_quizzes_for_test = 10 -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, - ) +if args.test == "tsne": + model = models[0] -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, - ) + quizzes = [] + labels = [] + nb_samples_per_task = 1000 -elif args.task == "stack": - task = TaskStack( - nb_train_samples=args.nb_train_samples, - nb_test_samples=args.nb_test_samples, - batch_size=args.batch_size, - nb_steps=args.stack_nb_steps, - nb_stacks=args.stack_nb_stacks, - nb_digits=args.stack_nb_digits, - fraction_values_for_train=args.stack_fraction_values_for_train, - device=device, - ) + for n, t in enumerate(args.grids_world_tasks.split(",")): + quizzes.append( + quiz_machine.problem.generate_w_quizzes(nb_samples_per_task, [t]) + ) + labels.append(torch.full((quizzes[-1].size(0),), n)) -elif args.task == "expr": - task = TaskExpr( - nb_train_samples=args.nb_train_samples, - nb_test_samples=args.nb_test_samples, - nb_variables=args.expr_nb_variables, - sequence_length=args.expr_sequence_length, - batch_size=args.batch_size, - device=device, - ) + quizzes = torch.cat(quizzes, dim=0) + labels = torch.cat(labels, dim=0) -else: - raise ValueError(f"Unknown task {args.task}") + with torch.autograd.no_grad(): + model.eval().to(main_device) + record = [] + for input, targets in zip( + quizzes.split(args.batch_size), labels.split(args.batch_size) + ): + input = input.to(main_device) + bs = mygpt.BracketedSequence(input) + bs = mygpt.BracketedSequence(F.pad(bs.x, (1, -1)), bs.first, bs.nb) + bs = model.embedding(bs) + bs = model.trunk[args.nb_blocks // 2](bs) + record.append((bs.x.to("cpu"), targets)) -###################################################################### + x = torch.cat([x for x, y in record], dim=0).flatten(1) + y = torch.cat([y for x, y in record], dim=0) -log_string(f"device {device}") + print(f"{x.size()=} {y.size()=}") + # torch.save((x,y), "/tmp/embed.pth") + # exit(0) -vocabulary_size = task.vocabulary_size() + from sklearn.manifold import TSNE -log_string(f"vocabulary_size {vocabulary_size}") + x_np = x.numpy() + z_np = TSNE(n_components=2, perplexity=50).fit_transform(x_np) + z = torch.from_numpy(z_np) -############################## - -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, -) + print(f"{z.size()=}") -model.to(device) + with open("/tmp/result.dat", "w") as f: + for k in range(z.size(0)): + f.write(f"{y[k]} {z[k,0]} {z[k,1]}\n") -nb_parameters = sum(p.numel() for p in model.parameters()) -log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)") + exit(0) ###################################################################### -nb_epochs_finished = 0 +if args.test == "generator": + token_prolog_0 = vocabulary_size + 0 + token_prolog_1 = vocabulary_size + 1 + token_prolog_2 = vocabulary_size + 2 + generator_vocabulary_size = vocabulary_size + 3 -if args.no_checkpoint: - log_string(f"not trying to load checkpoint.") + generator = mygpt.MyGPT( + vocabulary_size=generator_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, + compute_attzero=compute_causal_attzero, + dropout=args.dropout, + ).to(main_device) -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"]) + generator.main_test_accuracy = 0.0 - log_string(f"checkpoint loaded with {nb_epochs_finished} epochs finished.") + filename = f"generator.pth" + try: + d = torch.load(os.path.join(args.result_dir, filename)) + generator.load_state_dict(d[0]) + generator.main_test_accuracy = d[1] + log_string(f"successfully loaded {filename}") except FileNotFoundError: - log_string("starting from scratch.") + log_string(f"cannot find {filename}") + pass - except: - log_string("error when loading the checkpoint.") - exit(1) + for n_epoch in range(args.nb_epochs): + one_generator_epoch( + generator, + quiz_machine=quiz_machine, + models=models, + fraction_w_quizzes=1 if n_epoch < 25 else 0.5, + local_device=main_device, + ) -###################################################################### + filename = f"generator.pth" + torch.save( + (generator.state_dict(), generator.main_test_accuracy), + os.path.join(args.result_dir, filename), + ) + log_string(f"wrote {filename}") -nb_epochs = args.nb_epochs if args.nb_epochs > 0 else nb_epochs_default + c_quizzes, prologs = generate_c_quizzes_with_generator( + generator, quiz_machine, args.batch_size + ) -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) + seq_logproba = quiz_machine.models_logprobas( + models, c_quizzes, ("A", "f_A", "B", "f_B"), (0, 0, 0, 1), (0, 0, 1, 0) + ) + quiz_machine.models_logprobas( + models, c_quizzes, ("f_A", "A", "f_B", "B"), (0, 0, 0, 1), (0, 0, 1, 0) + ) -############################## + probas = seq_logproba.exp() -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(",") - ] - } - - 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 + u0 = probas <= args.proba_not_understands + u2 = probas >= args.proba_understands + u1 = (u0 | u2) == False -log_string(f"learning_rate_schedule {learning_rate_schedule}") + predicted_prologs = ( + (u0.long() * token_prolog_0) + + (u1.long() * token_prolog_1) + + (u2.long() * token_prolog_2) + ) -############################## + comments = [] -nb_samples_seen = 0 + nb_errors = (predicted_prologs != prologs).long().sum() + nb_total = prologs.numel() -if nb_epochs_finished >= nb_epochs: - task.produce_results(nb_epochs_finished, model) + log_string(f"generator_error {nb_errors} / {nb_total}") -for n_epoch in range(nb_epochs_finished, nb_epochs): - learning_rate = learning_rate_schedule[n_epoch] + def readable(prologs): + return (prologs == token_prolog_1) + 2 * (prologs == token_prolog_2) - log_string(f"learning_rate {learning_rate}") + for aa, ee, ff in zip(probas, readable(predicted_prologs), readable(prologs)): + sa = "prolog " + " ".join( + [f"{e.item()}/{f.item()}" for e, f in zip(ee, ff)] + ) + sp = "proba " + " ".join([f"{p.item():.02f}" for p in aa]) + comments.append(sa + "\n" + sp) - 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}.") + filename = f"generator_batch_{n_epoch:04d}.png" + quiz_machine.problem.save_quizzes_as_image( + args.result_dir, filename, c_quizzes, comments=comments + ) + log_string(f"wrote {filename}") - model.train() + exit(0) - nb_train_samples, acc_train_loss = 0, 0.0 +###################################################################### - 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) +for n_epoch in range(current_epoch, args.nb_epochs): + state = {"current_epoch": n_epoch} + filename = "state.pth" + torch.save(state, os.path.join(args.result_dir, filename)) + log_string(f"wrote {filename}") - optimizer.zero_grad() - loss.backward() - optimizer.step() + log_string(f"--- epoch {n_epoch} ----------------------------------------") - with torch.autograd.no_grad(): - model.eval() + cta = " ".join([f"{float(m.main_test_accuracy):.04f}" for m in models]) + log_string(f"current_test_accuracies {cta}") - nb_test_samples, acc_test_loss = 0, 0.0 + ################################################## + # If all the models are good enough, generate new quizzes and + # re-compute the test errors - for input in task.batches(split="test"): - input = input.to(device) + if min([m.main_test_accuracy for m in models]) >= args.accuracy_to_make_c_quizzes: + record_new_c_quizzes( + models, + quiz_machine, + nb_for_train=args.nb_new_c_quizzes_for_train, + nb_for_test=args.nb_new_c_quizzes_for_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) + filename = "c_quizzes.pth" + quiz_machine.save_c_quizzes(os.path.join(args.result_dir, filename)) + log_string(f"wrote {filename}") - train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples)) - test_perplexity = math.exp(min(100, acc_test_loss / nb_test_samples)) + # Force one epoch of training + for model in models: + model.main_test_accuracy = 0.0 - log_string( - f"perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}" + ################################################## + # Select, improve, and eval the worst model(s) + + ranked_models = sorted(models, key=lambda m: float(m.main_test_accuracy)) + + weakest_models = ranked_models[: len(gpus)] + + threads = [] + + for gpu, model in zip(gpus, weakest_models): + log_string(f"training model {model.id}") + + t = threading.Thread( + target=one_epoch, daemon=True, args=(model, quiz_machine, gpu) ) - task.produce_results(n_epoch, model) + threads.append(t) - checkpoint = { - "nb_epochs_finished": n_epoch + 1, - "model_state": model.state_dict(), - "rng_state": torch.get_rng_state(), - } + t.start() - if torch.cuda.is_available(): - checkpoint["cuda_rng_state"] = torch.cuda.get_rng_state() + for t in threads: + t.join() + + # Save the models to disk + + for model in weakest_models: + filename = f"gpt_{model.id:03d}.pth" + torch.save( + { + "state_dict": model.state_dict(), + "optimizer_state_dict": model.optimizer.state_dict(), + "main_test_accuracy": model.main_test_accuracy, + }, + os.path.join(args.result_dir, filename), + ) + log_string(f"wrote {filename}") + + for model in weakest_models: + save_additional_results(model, models, science_w_quizzes) + + ###################################################################### + + # Renew the training samples + + for model in weakest_models: + quiz_machine.renew_train_w_quizzes(model=model) - checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name) - torch.save(checkpoint, checkpoint_name) - log_string(f"saved checkpoint {checkpoint_name}") + if args.log_command is not None: + s = args.log_command.split() + s.insert(1, args.result_dir) + subprocess.run(s) ######################################################################