X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=world.py;h=61a07e9855bca6439c8c1e4b553a5b3ffced51ea;hb=910434daa74525b29d7ba117312fc67789d1ad84;hp=c4527f8835ad9a864cb07252c75ce6639d72b04f;hpb=6935899c1050d4f6a956fc8d2b50d2ba1544b6cc;p=picoclvr.git

diff --git a/world.py b/world.py
index c4527f8..61a07e9 100755
--- a/world.py
+++ b/world.py
@@ -146,7 +146,7 @@ def train_encoder(
 
         for input in tqdm.tqdm(train_input.split(batch_size), desc="vqae-train"):
             z = encoder(input)
-            zq = z if k < 1 else quantizer(z)
+            zq = z if k < 2 else quantizer(z)
             output = decoder(zq)
 
             output = output.reshape(
@@ -321,53 +321,6 @@ def generate_episode(steps, size=64):
 ######################################################################
 
 
-# ||x_i - c_j||^2 = ||x_i||^2 + ||c_j||^2 - 2<x_i, c_j>
-def sq2matrix(x, c):
-    nx = x.pow(2).sum(1)
-    nc = c.pow(2).sum(1)
-    return nx[:, None] + nc[None, :] - 2 * x @ c.t()
-
-
-def update_centroids(x, c, nb_min=1):
-    _, b = sq2matrix(x, c).min(1)
-    b.squeeze_()
-    nb_resets = 0
-
-    for k in range(0, c.size(0)):
-        i = b.eq(k).nonzero(as_tuple=False).squeeze()
-        if i.numel() >= nb_min:
-            c[k] = x.index_select(0, i).mean(0)
-        else:
-            n = torch.randint(x.size(0), (1,))
-            nb_resets += 1
-            c[k] = x[n]
-
-    return c, b, nb_resets
-
-
-def kmeans(x, nb_centroids, nb_min=1):
-    if x.size(0) < nb_centroids * nb_min:
-        print("Not enough points!")
-        exit(1)
-
-    c = x[torch.randperm(x.size(0))[:nb_centroids]]
-    t = torch.full((x.size(0),), -1)
-    n = 0
-
-    while True:
-        c, u, nb_resets = update_centroids(x, c, nb_min)
-        n = n + 1
-        nb_changes = (u - t).sign().abs().sum() + nb_resets
-        t = u
-        if nb_changes == 0:
-            break
-
-    return c, t
-
-
-######################################################################
-
-
 def generate_episodes(nb, steps):
     all_frames = []
     for n in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world-data"):
@@ -381,8 +334,6 @@ def create_data_and_processors(nb_train_samples, nb_test_samples, nb_epochs=10):
     train_input = generate_episodes(nb_train_samples, steps)
     test_input = generate_episodes(nb_test_samples, steps)
 
-    print(f"{train_input.size()=} {test_input.size()=}")
-
     encoder, quantizer, decoder = train_encoder(
         train_input, test_input, nb_epochs=nb_epochs
     )
@@ -394,26 +345,40 @@ def create_data_and_processors(nb_train_samples, nb_test_samples, nb_epochs=10):
     pow2 = (2 ** torch.arange(z.size(1), device=z.device))[None, None, :]
     z_h, z_w = z.size(2), z.size(3)
 
-    def frame2seq(x):
-        z = encoder(x)
-        ze_bool = (quantizer(z) >= 0).long()
-        seq = (
-            ze_bool.permute(0, 2, 3, 1).reshape(ze_bool.size(0), -1, ze_bool.size(1))
-            * pow2
-        ).sum(-1)
-        return seq
-
-    def seq2frame(seq, T=1e-2):
-        zd_bool = (seq[:, :, None] // pow2) % 2
-        zd_bool = zd_bool.reshape(zd_bool.size(0), z_h, z_w, -1).permute(0, 3, 1, 2)
-        logits = decoder(zd_bool * 2.0 - 1.0)
-        logits = logits.reshape(
-            logits.size(0), -1, 3, logits.size(2), logits.size(3)
-        ).permute(0, 2, 3, 4, 1)
-        results = torch.distributions.categorical.Categorical(
-            logits=logits / T
-        ).sample()
-        return results
+    def frame2seq(input, batch_size=25):
+        seq = []
+
+        for x in input.split(batch_size):
+            z = encoder(x)
+            ze_bool = (quantizer(z) >= 0).long()
+            output = (
+                ze_bool.permute(0, 2, 3, 1).reshape(
+                    ze_bool.size(0), -1, ze_bool.size(1)
+                )
+                * pow2
+            ).sum(-1)
+
+            seq.append(output)
+
+        return torch.cat(seq, dim=0)
+
+    def seq2frame(input, batch_size=25, T=1e-2):
+        frames = []
+
+        for seq in input.split(batch_size):
+            zd_bool = (seq[:, :, None] // pow2) % 2
+            zd_bool = zd_bool.reshape(zd_bool.size(0), z_h, z_w, -1).permute(0, 3, 1, 2)
+            logits = decoder(zd_bool * 2.0 - 1.0)
+            logits = logits.reshape(
+                logits.size(0), -1, 3, logits.size(2), logits.size(3)
+            ).permute(0, 2, 3, 4, 1)
+            output = torch.distributions.categorical.Categorical(
+                logits=logits / T
+            ).sample()
+
+            frames.append(output)
+
+        return torch.cat(frames, dim=0)
 
     return train_input, test_input, frame2seq, seq2frame
 
@@ -422,7 +387,10 @@ def create_data_and_processors(nb_train_samples, nb_test_samples, nb_epochs=10):
 
 if __name__ == "__main__":
     train_input, test_input, frame2seq, seq2frame = create_data_and_processors(
-        10000, 1000
+        # 10000, 1000,
+        100,
+        100,
+        nb_epochs=2,
     )
 
     input = test_input[:64]