Update.

[picoclvr.git] / tasks.py

diff --git a/tasks.py b/tasks.py
index b33dee2..f4be293 100755 (executable)
--- a/tasks.py
+++ b/tasks.py
@@ -111,13 +111,25 @@ 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(self.train_ar_mask.unique()) == (0, 1)
-            and tuple(self.test_ar_mask.unique()) == (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"}
@@ -151,17 +163,24 @@ 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 = ar_mask.sum().item()
-            nb_correct = ((result == input).long() * ar_mask).sum().item()
+            # nb_total = ar_mask.sum().item()
+            # nb_correct = ((result == input).long() * ar_mask).sum().item()
 
             return nb_total, nb_correct
 
@@ -193,30 +212,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}")
 
 
 ######################################################################