Update.

[picoclvr.git] / tasks.py

diff --git a/tasks.py b/tasks.py
index c0ad5ff..443419e 100755 (executable)
--- a/tasks.py
+++ b/tasks.py
@@ -754,15 +754,17 @@ class Maze(Task):
     def compute_error(
         self, model, split="train", nb_to_use=-1, deterministic_synthesis=False
     ):
+        model_device = next(model.parameters()).device
         nb_total, nb_correct = 0, 0
         count = torch.zeros(
             self.width * self.height,
             self.width * self.height,
-            device=self.device,
+            device=model_device,
             dtype=torch.int64,
         )
 
         for input in self.batches(split, nb_to_use):
+            input = input.to(model_device)
             result = input.clone()
             ar_mask = result.new_zeros(result.size())
             ar_mask[:, self.height * self.width :] = 1
@@ -836,7 +838,7 @@ class Maze(Task):
                         eol = " " if j < count.size(1) - 1 else "\n"
                         f.write(f"{count[i,j]}{eol}")
 
-        input = self.test_input[:48]
+        input = self.test_input[:48].to(next(model.parameters()).device)
         result = input.clone()
         ar_mask = result.new_zeros(result.size())
         ar_mask[:, self.height * self.width :] = 1
@@ -1098,6 +1100,34 @@ class Stack(Task):
             device=self.device,
         )
 
+        #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+        for label, input in [
+            ("train", self.train_input[:32]),
+            ("test", self.test_input[:32]),
+        ]:
+            output = model(BracketedSequence(input)).x
+            output = output.log_softmax(dim=-1)
+            filename = os.path.join(
+                result_dir, f"stack_with_crossentropy_{n_epoch:04d}_{label}.txt"
+            )
+            with open(filename, "w") as f:
+                for n in range(input.size(0)):
+                    s = stack.seq_to_str(
+                        input[n], nb_stacks=self.nb_stacks, nb_digits=self.nb_digits
+                    )
+                    for t, k, w in zip(range(input[n].size(0)), input[n], s.split(" ")):
+                        u = (
+                            " " * (10 - len(w))
+                            + w
+                            + " "
+                            + str(output[n][t][k].exp().item())
+                            + "\n"
+                        )
+                        f.write(u)
+                    f.write("\n")
+            logger(f"wrote {filename}")
+        #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
         for n in range(result.size(0)):
             logger(
                 f"test_after  {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}"