X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=08aa8caf997a2b54ca3cea8fb29dd784c18820e8;hb=128d372813e99d8474bb6e967d5c7e7f085c819d;hp=0858282fa3b5cab185d14025d5ad758de44411c2;hpb=0f4c86c0e7730db4147f136df5aeb5528fc943a0;p=picoclvr.git diff --git a/tasks.py b/tasks.py index 0858282..08aa8ca 100755 --- a/tasks.py +++ b/tasks.py @@ -14,10 +14,8 @@ from torch.nn import functional as F from mygpt import BracketedSequence -try: - from graph import save_attention_image -except ImportError: - save_attention_image = None +# from graph import save_attention_image +save_attention_image = None ###################################################################### @@ -110,15 +108,26 @@ class SandBox(Task): self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1 - # A bit of paranoia never hurts - assert ( - self.nb_codes <= max_nb_codes - and self.train_input.min() >= 0 - and self.test_input.min() >= 0 - and tuple(x.item() for x in self.train_ar_mask.unique()) in { (0,), (1,), (0,1) } - and tuple(x.item() for x in self.test_ar_mask.unique()) in { (0,), (1,), (0,1) } - ) + assert self.nb_codes <= max_nb_codes + assert self.train_input.min() >= 0 + assert self.test_input.min() >= 0 + assert tuple(x.item() for x in self.train_ar_mask.unique()) in { + (0,), + (1,), + (0, 1), + } + assert tuple(x.item() for x in self.test_ar_mask.unique()) in { + (0,), + (1,), + (0, 1), + } + + if logger is not None: + for s, a in zip(self.train_input[:100], self.train_ar_mask[:100]): + logger(f"train_sequences {self.problem.seq2str(s)}") + a = "".join(["01"[x.item()] for x in a]) + logger(f" {a}") def batches(self, split="train", nb_to_use=-1, desc=None): assert split in {"train", "test"} @@ -152,16 +161,21 @@ class SandBox(Task): device=self.device, ) + log_ground_truth = ar_mask.min() == 0 + if logger is not None: for sp, st in zip(result[:10], input[:10]): logger( f"test_sequences {n_epoch} prediction {self.problem.seq2str(sp)}" ) - logger( - f" {n_epoch} ground truth {self.problem.seq2str(st)}" - ) + if log_ground_truth: + logger( + f" {n_epoch} ground truth {self.problem.seq2str(st)}" + ) - nb_total, nb_correct = self.problem.compute_nb_correct(input, ar_mask, result) + nb_total, nb_correct = self.problem.compute_nb_correct( + input, ar_mask, result + ) # nb_total = ar_mask.sum().item() # nb_correct = ((result == input).long() * ar_mask).sum().item() @@ -186,9 +200,7 @@ class SandBox(Task): logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}") - if save_attention_image is None: - logger("no save_attention_image (is pycairo installed?)") - else: + if save_attention_image is not None: for k in range(10): ns = torch.randint(self.test_input.size(0), (1,)).item() input = self.test_input[ns : ns + 1].clone() @@ -196,30 +208,30 @@ class SandBox(Task): with torch.autograd.no_grad(): t = model.training model.eval() - model.record_attention(True) + # model.record_attention(True) model(BracketedSequence(input)) model.train(t) - ram = model.retrieve_attention() - model.record_attention(False) - - tokens_output = [c for c in self.problem.seq2str(input[0])] - tokens_input = ["n/a"] + tokens_output[:-1] - for n_head in range(ram[0].size(1)): - filename = os.path.join( - result_dir, f"sandbox_attention_{k}_h{n_head}.pdf" - ) - attention_matrices = [m[0, n_head] for m in ram] - save_attention_image( - filename, - tokens_input, - tokens_output, - attention_matrices, - k_top=10, - # min_total_attention=0.9, - token_gap=12, - layer_gap=50, - ) - logger(f"wrote {filename}") + # ram = model.retrieve_attention() + # model.record_attention(False) + + # tokens_output = [c for c in self.problem.seq2str(input[0])] + # tokens_input = ["n/a"] + tokens_output[:-1] + # for n_head in range(ram[0].size(1)): + # filename = os.path.join( + # result_dir, f"sandbox_attention_{k}_h{n_head}.pdf" + # ) + # attention_matrices = [m[0, n_head] for m in ram] + # save_attention_image( + # filename, + # tokens_input, + # tokens_output, + # attention_matrices, + # k_top=10, + ##min_total_attention=0.9, + # token_gap=12, + # layer_gap=50, + # ) + # logger(f"wrote {filename}") ###################################################################### @@ -1463,6 +1475,7 @@ class Grid(Task): nb_test_samples, batch_size, size, + fraction_play=0.0, logger=None, device=torch.device("cpu"), ): @@ -1478,10 +1491,12 @@ class Grid(Task): ) self.train_descr = self.grid_factory.generate_samples( - nb_train_samples, lambda r: tqdm.tqdm(r) + nb=nb_train_samples, + fraction_play=fraction_play, + progress_bar=lambda r: tqdm.tqdm(r), ) self.test_descr = self.grid_factory.generate_samples( - nb_test_samples, lambda r: tqdm.tqdm(r) + nb=nb_test_samples, fraction_play=0.0, progress_bar=lambda r: tqdm.tqdm(r) ) # Build the tokenizer @@ -1551,6 +1566,42 @@ class Grid(Task): logger(f"test_performance {n_epoch} {nb_total=} {nb_correct=}") logger(f"main_test_accuracy {n_epoch} {nb_correct / nb_total}") + if n_epoch == 5 or n_epoch == 10 or n_epoch == 20: + if save_attention_image is None: + logger("no save_attention_image (is pycairo installed?)") + else: + for k in range(10): + ns = k # torch.randint(self.test_input.size(0), (1,)).item() + input = self.test_input[ns : ns + 1].clone() + with torch.autograd.no_grad(): + t = model.training + model.eval() + model.record_attention(True) + model(BracketedSequence(input)) + model.train(t) + ram = model.retrieve_attention() + model.record_attention(False) + + tokens_output = [self.id2token[t.item()] for t in input[0]] + tokens_input = ["n/a"] + tokens_output[:-1] + for n_head in range(ram[0].size(1)): + filename = os.path.join( + result_dir, + f"sandbox_attention_epoch_{n_epoch}_sample_{k}_head_{n_head}.pdf", + ) + attention_matrices = [m[0, n_head] for m in ram] + save_attention_image( + filename, + tokens_input, + tokens_output, + attention_matrices, + k_top=10, + # min_total_attention=0.9, + token_gap=12, + layer_gap=50, + ) + logger(f"wrote {filename}") + ######################################################################