X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=tasks.py;h=b967465de6c23a94b9ee1af8a672337167903ac8;hb=8a548630c88957264306db4354e880414b0fa8ef;hp=50d541b11ccab4fb1dca97815412e1e88d69fc93;hpb=9047bd8185ed99c1302d8812551af3d5bd4602cb;p=culture.git diff --git a/tasks.py b/tasks.py index 50d541b..b967465 100755 --- a/tasks.py +++ b/tasks.py @@ -14,9 +14,6 @@ from torch.nn import functional as F from mygpt import BracketedSequence -# from graph import save_attention_image -save_attention_image = None - ###################################################################### @@ -25,6 +22,7 @@ def masked_inplace_autoregression( batch_size, input, ar_mask, + temperature, deterministic_synthesis, forbidden_tokens=None, logit_biases=None, @@ -47,17 +45,22 @@ def masked_inplace_autoregression( t = model.training model.eval() + sum_logits = 0 + for input, ar_mask in batches: - model.masked_inplace_autoregression( - input, - ar_mask, - deterministic_synthesis, - forbidden_tokens, - logit_biases, + sum_logits += model.masked_inplace_autoregression( + input=input, + ar_mask=ar_mask, + temperature=temperature, + deterministic_synthesis=deterministic_synthesis, + forbidden_tokens=forbidden_tokens, + forced_biases=logit_biases, ) model.train(t) + return sum_logits + ###################################################################### @@ -82,9 +85,9 @@ import world class World(Task): def save_image(self, input, result_dir, filename, logger): - img = world.sample2img(input.to("cpu"), self.height, self.width) + img = world.seq2img(input.to("cpu"), self.height, self.width) image_name = os.path.join(result_dir, filename) - torchvision.utils.save_image(img.float() / 255.0, image_name, nrow=8, padding=2) + torchvision.utils.save_image(img.float() / 255.0, image_name, nrow=6, padding=4) logger(f"wrote {image_name}") def make_ar_mask(self, input): @@ -107,11 +110,11 @@ class World(Task): self.height = 6 self.width = 8 - self.train_input = world.generate( + self.train_input = world.generate_seq( nb_train_samples, height=self.height, width=self.width ).to(device) - self.test_input = world.generate( + self.test_input = world.generate_seq( nb_test_samples, height=self.height, width=self.width ).to(device) @@ -122,7 +125,7 @@ class World(Task): if result_dir is not None: self.save_image( - self.train_input[:96], result_dir, f"world_train.png", logger + self.train_input[:72], result_dir, f"world_train.png", logger ) def batches(self, split="train", desc=None): @@ -170,11 +173,12 @@ class World(Task): result = input.clone() * (1 - ar_mask) masked_inplace_autoregression( - model, - self.batch_size, - result, - ar_mask, - deterministic_synthesis, + model=model, + batch_size=self.batch_size, + input=result, + ar_mask=ar_mask, + temperature=1.0, + deterministic_synthesis=deterministic_synthesis, progress_bar_desc=None, device=self.device, ) @@ -208,24 +212,33 @@ class World(Task): result = input.clone() * (1 - ar_mask) masked_inplace_autoregression( - model, - self.batch_size, - result, - ar_mask, - deterministic_synthesis, + model=model, + batch_size=self.batch_size, + input=result, + ar_mask=ar_mask, + temperature=1.0, + deterministic_synthesis=deterministic_synthesis, progress_bar_desc=None, device=self.device, ) self.save_image( - result[:96], + result[:72], result_dir, - f"world_result_{n_epoch:04d}_{model.id:02d}.png", + f"world_prediction_{n_epoch:04d}_{model.id:02d}.png", logger, ) return main_test_accuracy + def renew_samples(self, nb, for_train=True): + input = self.train_input if for_train else self.test_input + nb = min(nb, input.size(0)) + input[:-nb] = input[nb:].clone() + input[-nb:] = world.generate_seq(nb, height=self.height, width=self.width).to( + self.device + ) + def store_new_quizzes(self, new_quizzes, for_train=True): if for_train: self.train_quizzes.append(new_quizzes) @@ -240,67 +253,105 @@ class World(Task): nb, model, other_models, + desired_average_logits=None, ): - new_quizzes = torch.empty( + ############################################################### + # Generate quizzes with model + + quizzes = torch.empty( nb, self.height * self.width * 2 + 1, device=self.device, dtype=torch.int64 ) - ar_mask = torch.full(new_quizzes.size(), 1, device=self.device) - - masked_inplace_autoregression( - model, - self.batch_size, - new_quizzes, - ar_mask, + ar_mask = torch.full(quizzes.size(), 1, device=self.device) + + sum_logits = masked_inplace_autoregression( + model=model, + batch_size=self.batch_size, + input=quizzes, + ar_mask=ar_mask, + temperature=1.0, deterministic_synthesis=False, - progress_bar_desc="new quizzes", + progress_bar_desc="creating quizzes", device=self.device, ) - ar_mask = self.make_ar_mask(new_quizzes) + average_logits = sum_logits / quizzes.numel() + + # It's a bit brutal to do it twice, we should probably have a + # moving average and apply it right away + + if desired_average_logits is not None: + temperature = average_logits / desired_average_logits + masked_inplace_autoregression( + model=model, + batch_size=self.batch_size, + input=quizzes, + ar_mask=ar_mask, + temperature=temperature, + deterministic_synthesis=False, + progress_bar_desc="creating quizzes", + device=self.device, + ) + + ############################################################### + # Create the reverse quizzes + + l = self.height * self.width + direction = quizzes[:, l : l + 1] + direction = world.token_forward * ( + direction == world.token_backward + ) + world.token_backward * (direction == world.token_forward) + reverse_quizzes = torch.cat( + [quizzes[:, l + 1 :], direction, quizzes[:, :l]], dim=1 + ) + + ar_mask = self.make_ar_mask(quizzes) - nb_correct = 0 + ############################################################### + # Check how many of the other models can solve them in both + # directions + + nb_correct = [] for m in other_models: - result = new_quizzes.clone() + result = quizzes.clone() masked_inplace_autoregression( - m, - self.batch_size, - result, - ar_mask, + model=m, + batch_size=self.batch_size, + input=result, + ar_mask=ar_mask, + temperature=1.0, deterministic_synthesis=True, progress_bar_desc="solving quizzes", device=self.device, ) - l = self.height * self.width - direction = new_quizzes[:, l : l + 1] - direction = world.token_forward * ( - direction == world.token_backward - ) + world.token_backward * (direction == world.token_forward) - inverted_quizzes = torch.cat( - [new_quizzes[:, l + 1 :], direction, new_quizzes[:, :l]], dim=1 - ) + correct = (quizzes == result).long().min(dim=-1).values - inverted_result = inverted_quizzes.clone() + reverse_result = reverse_quizzes.clone() masked_inplace_autoregression( - m, - self.batch_size, - inverted_result, - ar_mask, + model=m, + batch_size=self.batch_size, + input=reverse_result, + ar_mask=ar_mask, + temperature=1.0, deterministic_synthesis=True, - progress_bar_desc="solving reverse quizzes", + progress_bar_desc="solving reversed quizzes", device=self.device, ) - nb_correct += ( - ( - (new_quizzes == result).long() - * (inverted_quizzes, inverted_result).long() - ) - .min(dim=-1) - .values + reverse_correct = ( + (reverse_quizzes == reverse_result).long().min(dim=-1).values ) - return new_quizzes, nb_correct + nb_correct.append((correct * reverse_correct)[None, :]) + + nb_correct = torch.cat(nb_correct, dim=0) + + filename = os.path.join(result_dir, "correct_{n_epoch:04d}.dat") + with open(filename, "w") as f: + for k in nb_correct: + f.write(f"{k}\n") + + return quizzes, nb_correct.sum(dim=0), average_logits