Update.

[picoclvr.git] / tasks.py

diff --git a/tasks.py b/tasks.py
index 82d965b..96d0621 100755 (executable)
--- a/tasks.py
+++ b/tasks.py
@@ -20,6 +20,8 @@ def masked_inplace_autoregression(
     progress_bar_desc="autoregression",
     device=torch.device("cpu"),
 ):
+    assert input.size() == ar_mask.size()
+
     batches = zip(input.split(batch_size), ar_mask.split(batch_size))
 
     if progress_bar_desc is not None:
@@ -27,7 +29,7 @@ def masked_inplace_autoregression(
             batches,
             dynamic_ncols=True,
             desc=progress_bar_desc,
-            total=input.size(0) // batch_size,
+            #total=input.size(0) // batch_size,
         )
 
     with torch.autograd.no_grad():
@@ -590,8 +592,6 @@ class Snake(Task):
             )
             result *= 1 - ar_mask
 
-            # snake.solver(result,ar_mask)
-
             masked_inplace_autoregression(
                 model,
                 self.batch_size,
@@ -605,19 +605,8 @@ class Snake(Task):
 
             nb_correct = ((result == input).long() * (prior_visits > 0) * ar_mask).sum()
 
-            # nb_total = result.size(0)
-            # nb_correct = ((result - input).abs().sum(1) == 0).sum()
-
             return nb_total, nb_correct
 
-        # train_nb_total, train_nb_correct = compute_nb_correct(
-        # self.train_input, self.train_prior_visits
-        # )
-
-        # logger(
-        # f"accuracy_train nb_total {train_nb_total} nb_correct {train_nb_correct} accuracy {(100.0*train_nb_correct)/train_nb_total:.02f}%"
-        # )
-
         test_nb_total, test_nb_correct = compute_nb_correct(
             self.test_input[:1000], self.test_prior_visits[:1000]
         )
@@ -880,17 +869,22 @@ class Expr(Task):
             values_input = expr.extract_results([self.seq2str(s) for s in input])
             values_result = expr.extract_results([self.seq2str(s) for s in result])
 
-            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):
+            filename = os.path.join(result_dir, f"expr_result_{n_epoch:04d}.txt")
+
+            with open(filename, "w") as f:
+                for i, r in zip(values_input, values_result):
+                    for n, vi in i.items():
+                        vr = r.get(n)
+                        f.write(f"{vi} {-1 if vr is None else vr}\n")
+
+                        if vr is None or vr < 0:
                             nb_missed += 1
                         else:
-                            nb_delta[d] += 1
+                            d = abs(vr - vi)
+                            if d >= nb_delta.size(0):
+                                nb_missed += 1
+                            else:
+                                nb_delta[d] += 1
 
             ######################################################################
 
@@ -952,3 +946,91 @@ class Expr(Task):
 
 
 ######################################################################
+
+import world
+
+
+class World(Task):
+    def __init__(
+        self,
+        nb_train_samples,
+        nb_test_samples,
+        batch_size,
+        vqae_nb_epochs,
+        device=torch.device("cpu"),
+    ):
+        self.batch_size = batch_size
+        self.device = device
+
+        (
+            train_frames,
+            self.train_actions,
+            test_frames,
+            self.test_actions,
+            self.frame2seq,
+            self.seq2frame,
+        ) = world.create_data_and_processors(
+            nb_train_samples,
+            nb_test_samples,
+            mode="first_last",
+            nb_steps=30,
+            nb_epochs=vqae_nb_epochs,
+            device=device,
+        )
+
+        self.train_input = self.frame2seq(train_frames)
+        self.train_input = self.train_input.reshape(self.train_input.size(0) // 2, -1)
+        self.test_input = self.frame2seq(test_frames)
+        self.test_input = self.test_input.reshape(self.test_input.size(0) // 2, -1)
+
+        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, result_dir, logger, deterministic_synthesis
+    ):
+        l = self.train_input.size(1)
+        k = torch.arange(l, device=self.device)[None, :]
+        result = self.test_input[:64].clone()
+
+        ar_mask = (k >= l // 2).long().expand_as(result)
+        result *= 1 - ar_mask
+
+        masked_inplace_autoregression(
+            model,
+            self.batch_size,
+            result,
+            ar_mask,
+            deterministic_synthesis,
+            device=self.device,
+        )
+
+        result = result.reshape(result.size(0) * 2, -1)
+
+        frames = self.seq2frame(result)
+        image_name = os.path.join(result_dir, f"world_result_{n_epoch:04d}.png")
+        torchvision.utils.save_image(
+            frames.float() / (world.Box.nb_rgb_levels - 1),
+            image_name,
+            nrow=8,
+            padding=1,
+            pad_value=0.0,
+        )
+        logger(f"wrote {image_name}")
+
+
+######################################################################