model exploration now uses ConvVAE

Author: soof-golan

src/explore_model.py

@@ -18,20 +18,33 @@
 
 import matplotlib.pyplot as plt
 import torch
+import torchvision.transforms.functional as TF  # noqa: N812
+from einops import rearrange
 from IPython.core.display_functions import display
 from ipywidgets import interact
 from torchvision.transforms import ToTensor
 from torchvision.transforms.functional import to_pil_image
 
+import models
 from datamodules import ImagesDataModule
-from models import FullyConnectedAutoEncoder
+from train import LOGS_DIR
 
 # %%
+DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+DEVICE = torch.device("mps") if torch.backends.mps.is_available() else DEVICE
+
+# %%
+ModelClass = models.ConvVAE
+dataset_name = "FashionMNIST"
+datamodule = ImagesDataModule(dataset_name, 1, 10)
+
+# %%
+model_name = ModelClass.__name__.lower()
 ckpt_dir = (
-    Path("/tmp/logs")
-    / "fullyconnectedautoencodersgd-fashionmnist"
-    / "fullyconnectedautoencodersgd-fashionmnist"
+    LOGS_DIR
+    / f"{model_name}-{dataset_name.lower()}/{model_name}-{dataset_name.lower()}"
 )
+
 for p in ckpt_dir.parents[::-1] + (ckpt_dir,):
     if not p.exists():
         raise ValueError(f"{p} not exists")
@@ -53,54 +66,85 @@ def sort_dict(d: dict):
 all_ckpts = sort_dict(get_last_fn(subdir) for subdir in ckpt_dir.glob("*"))
 display(all_ckpts)
 
+
 # %%
 # torch.load(ckpt_dir/list(all_ckpts.values())[-1])['hyper_parameters']
 
 # %%
-model = FullyConnectedAutoEncoder.load_latest_checkpoint(ckpt_dir)
-model.eval()
+
+
+def load_model():
+    return ModelClass.load_latest_checkpoint(ckpt_dir, map_location=DEVICE).eval()
+
+
+model = load_model()
 print(model.hparams)
 print(model)
 
 # %%
-x_rand = torch.rand(1, 1, 28, 28)
-image = ImagesDataModule("FashionMNIST", 1, 10).dataset()[0][0]
+x_rand = torch.rand(1, 1, 32, 32)
+image, _target = datamodule.dataset()[0]
 
 x_real = ToTensor()(image).unsqueeze(0)
+x_rand = TF.center_crop(x_rand, 32)
+x_real = TF.center_crop(x_real, 32)
 print(x_real.shape)
 
 
 # %%
-def show_tensors(imgs: list[torch.Tensor]):
+def show_tensors(imgs: list[torch.Tensor], normalize=True, figsize=None):
     if not isinstance(imgs, list):
         imgs = [imgs]
-    fig, axss = plt.subplots(ncols=len(imgs), squeeze=False)
+    fig, axss = plt.subplots(ncols=len(imgs), squeeze=False, figsize=figsize)
     axs = axss[0]
     for i, img in enumerate(imgs):
-        img_clipped = img.detach().clip(0, 1)
-        img_pil = to_pil_image(img_clipped)
+        if normalize:
+            img = (img - img.min()) / (img.max() - img.min())
+        img = img.clamp(0, 1).detach()
+        img_pil = to_pil_image(img)
         axs[i].imshow(img_pil, cmap="gray", vmin=0, vmax=255)
         axs[i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
 
 
 for x in [x_rand, x_real]:
-    show_tensors([x[0], model(x.cuda())[0]])
+    show_tensors([x[0], model(x.to(DEVICE)).x_hat[0]])
 
 # %%
-n_latent = 8
+n_latent = model.latent_dim
 
-lims = (-2, 2, 0.01)
+lims = (-3, 3, 0.01)
 all_lims = {f"x{i:02}": lims for i in range(n_latent)}
 
 
 def show_from_latent(**inputs):
     data = torch.tensor(list(inputs.values()))
-    data = data.view(1, -1).cuda()
+    data = data.view(1, -1).to(DEVICE)
     result = model.decoder(data)[0]
-    show_tensors(result)
+    show_tensors(result, normalize=True)
     plt.show()
 
 
 interact(show_from_latent, **all_lims)
 
 # %%
+model = load_model()
+
+
+def sample_latent(model, n: int = 30, lim: float = 3.0, downsample_factor: int = 2):
+    x = torch.linspace(-lim, lim, n)
+    y = torch.linspace(-lim, lim, n)
+    z = torch.cartesian_prod(x, y)
+    assert z.shape[1] == 2
+    with torch.inference_mode():
+        outs = model.decoder(z.to(model.device))
+        out = rearrange(outs, "(i j) c h w -> c (i h) (j w)", i=n, j=n)
+        out = torch.nn.functional.avg_pool2d(out, kernel_size=downsample_factor)
+        # out = reduce(out, "c (h i) (w j) -> c h w", i=downsample_factor,j=downsample_factor, reduction="max")
+    return out
+
+
+out = sample_latent(model)
+print(out.shape)
+show_tensors(out, figsize=(10, 10))
+
+# %%