Try VAE with additional AE reconstruction loss

README.md

@@ -19,44 +19,12 @@ Then, you can import the module `my_pytorch_kit`.
 ## Usage
 
 This package revolves around the [`BaseModel`](my_pytorch_kit/model/models.py), [`Trainer`](my_pytorch_kit/train/train.py) and [`Evaluator`](my_pytorch_kit/evaluation/evaluation.py) classes, which are extended to model, train and evaluate a model respectively.  
-```mermaid
-graph TD
-    subgraph "Usage Workflow"
-        A["<b>Define Model</b><br/>(extends BaseModel)"]
-        B["<b>Define Data</b><br/>(Dataset / DataLoader)"]
-        G["<b>Extras</b><br/>(Optimizer, Tensorboard)"]
-        C["<b>Initialize Trainer</b><br/>(extends Trainer)"]
-        D["<b>Initialize Evaluator</b><br/>(extends Evaluator)"]
-        E["trainer.train()"]
-        F["evaluator.evaluate()"]
-        G["<b>Intitialize Tuner</b><br/>(Hyperparameter Tuning)"]
-        H["tuner.tune()"]
-        I["<b>Provided Architectures</b><br/>(Classifier, AE, VAE, ...)"]
 
-    end
 
-    %% Define node relationships
-    A --> C
-    B --> C
-    I --> A
-    A --> D
-    B --> D
-    C --> E
-    D --> F
-    C --> G
-    G --> H
+<p align="middle">
+  <img src="docs/assets/flowchart.svg" width="600" />
+</p>
 
-    %% Style the nodes
-    style A fill:#fbe,stroke:#333,stroke-width:2px
-    style B fill:#fbe,stroke:#333,stroke-width:2px
-    style G fill:#ffc,stroke:#333,stroke-width:2px
-    style C fill:#cde,stroke:#333,stroke-width:2px
-    style D fill:#cde,stroke:#333,stroke-width:2px
-    style E fill:#cfc,stroke:#333,stroke-width:2px
-    style F fill:#cfc,stroke:#333,stroke-width:2px
-    style G fill:#cde,stroke:#333,stroke-width:2px
-    style H fill:#cfc,stroke:#333,stroke-width:2px
-```
 
 Furthermore, this package provides architecture implementations and modelling utilities.
 Currently implemented architectures include:

examples/mnist/classifier/classifier.py

@@ -58,7 +58,7 @@ def evaluate_batch(self, model, batch):
 
         trainer.train(optimizer, **hparams)
 
-        model.save_model("models/classifier.pt")
+        model.save_model("models/classifier.pt", hparams)
 
     else:
         model.load_model("models/classifier.pt")

examples/mnist/classifier/example.py

@@ -92,7 +92,7 @@ def accumulate_result(self, result):
     def get_result(self):
         return self.result
 
-    def on_eval(self):
+    def on_eval(self, model):
         self.batch_count = 0
 
 #++++++++++++++  Here comes the main code ++++++++++++++++#

examples/mnist/classifier/tune.py

@@ -85,7 +85,8 @@
                                              num_search=num_search,
                                              ranks_considered=ranks_considered,
                                              check_multiplicant=check_multiplier,
-                                             mode='dynamic')
+                                             mode='dynamic',
+                                             model_path='models/mnist_random_dynamic_')
 
     result = evaluator.evaluate(model, test_loader)
     print(f"Model accuracy on test set doing random dynamic search: {result*100:.2f}%")

examples/mnist/vae/vae.py

@@ -20,25 +20,26 @@ def kl_annealing(epoch, epochs, model, **kwargs):
     train = input("Train new model? (y/[n]) ").lower() == "y"
 
     hparams = {
-        "learning_rate": 1e-3,
+        "learning_rate": 3e-3,
         "batch_size": 64,
         "epochs": 15,
         "optimizer_method": "Adam",
         # "optimizer_kwargs": {"weight_decay": 1e-4},
         "loss_func": nn.BCELoss(reduction="sum"),
         "alpha": 1e-1,
-        "beta": 1.5e-3,
-        "classifier_loss_weight": 1.2,
+        "beta": 1.5e-1,
+        "classifier_loss_weight": 12,
         "start_beta": 1e-3,
         "final_beta": 4e-2,
         "feature_space": (64, 7, 7),
-        "latent_dim": 4,
+        "latent_dim": 2,
         "sample_input_shape": (1, 1, 28, 28),
         "classifier_num_layers": 3,
         # "epoch_function": kl_annealing
     }
 
     model = ImageVAESemiSupervised(**hparams)
+    model.proper_weight_init()
     # model = ImageVAE(**hparams)
 
     torchsummary.summary(model, (1, 28, 28))

examples/mnist/vae/vae_semi_generate.py

@@ -4,13 +4,25 @@
 import os
 
 from my_pytorch_kit.model.vae import ImageVAESemiSupervised
+from my_pytorch_kit.model.classifier import ImageClassifier
+from my_pytorch_kit.model.ae import ImageAE
 
-def generate_samples(model, latent_dim, num_samples=20):
+def generate_samples(model, latent_dim, num_samples=20, check_classifier=False, check_ae=False):
     """
     Generates num_samples samples from the model,
     plotting them in a grid with their labels
     """
 
+    classifier = None
+    if check_classifier:
+        classifier = ImageClassifier()
+        classifier.load_model("models/classifier.pt")
+
+    ae = None
+    if check_ae:
+        ae = ImageAE()
+        ae.load_model("models/ae.pt")
+
     num_samples = int(num_samples)
     n_rows = num_samples // 10
     n_cols = 10
@@ -20,13 +32,40 @@ def generate_samples(model, latent_dim, num_samples=20):
         z = torch.randn((num_samples, latent_dim))
         images, labels = model.generate(z)
 
+        cl_labels = None
+        if classifier:
+            cl_labels = classifier(images)
+            cl_labels = torch.argmax(cl_labels, dim=1)
+
+        batch_recon_losses = None
+        if ae:
+            reconstruction = ae(images)
+            recon_losses = torch.nn.BCELoss(reduction="none")(reconstruction, images)
+            batch_recon_losses = recon_losses.mean(dim=(1, 2, 3))
+
         plt.figure(figsize=(10, 10))
         for i in range(num_samples):
             plt.subplot(n_rows, n_cols, i + 1)
             plt.imshow(images[i].view(28, 28).cpu().numpy(), cmap="Greys_r")
-            label = torch.argmax(labels[i]).item()
+            label = f"{torch.argmax(labels[i]).item()}"
+            if cl_labels is not None:
+                label += f" ({cl_labels[i].item()})"
+            if recon_losses is not None:
+                label += f" ({batch_recon_losses[i].item():.2f})"
+
             plt.title(label)
             plt.axis("off")
+
+        extra_text = "Generator label"
+        if cl_labels is not None:
+            extra_text += " (Classifier label)"
+        if recon_losses is not None:
+            extra_text += " (Reconstruction loss)"
+        plt.text(-100, 50, extra_text, ha="center", va="center")
+
+        print(labels.shape, cl_labels.shape)
+        print(torch.nn.CrossEntropyLoss()(labels, cl_labels).item())
+
         plt.show()
 
 
@@ -76,7 +115,7 @@ def generate_dataset(model, num_samples, latent_dim, std_divider=1.5, batch_size
 
     model.load_model("models/vae_semi_3.pt")
 
-    generate_samples(model, hparams["latent_dim"])
+    generate_samples(model, hparams["latent_dim"], check_classifier=True, check_ae=True)
 
     gen_dataset = input("Generate dataset? (y/[n]): ").lower() == "y"
 

examples/mnist/vae/vae_semi_tune.py

@@ -0,0 +1,148 @@
+
+import torch
+import torch.nn as nn
+import tqdm
+
+from my_pytorch_kit.model.vae import ImageVAESemiSupervised
+from my_pytorch_kit.train.tune import Tuner
+from my_pytorch_kit.evaluation.evaluation import Evaluator 
+from my_pytorch_kit.model.classifier import ImageClassifier
+from my_pytorch_kit.model.ae import ImageAE
+from mnist.utils.mnist_utils import get_mnist_loaders
+from mnist.autoencoder.ae_plots import plot_latent, plot_reconstructions, plot_label_clusters
+
+
+class GenerationEvaluator(Evaluator):
+
+
+    def __init__(self,
+                 latent_dim = 2,
+                 num_samples = 1000,
+                 batch_size = 64,
+                 ce_weight = 1,
+                 bce_weight = 2.5):
+        super().__init__()
+        self.batch_count = 0
+        self.result = 0
+
+        self.classifier = ImageClassifier()
+        self.classifier.load_model("models/classifier.pt")
+
+        self.ae = ImageAE()
+        self.ae.load_model("models/ae.pt")
+
+        self.latent_dim = latent_dim
+        self.num_samples = num_samples
+        self.batch_size = batch_size
+
+        self.bce = nn.BCELoss(reduction="mean")
+        self.ce = nn.CrossEntropyLoss(reduction="mean", label_smoothing=0.1)
+
+        self.ce_weight = ce_weight
+        self.bce_weight = bce_weight
+
+    def on_eval(self, model):
+
+        print("\nEvaluating...")
+        generated_count = 0
+
+        acc = 0
+
+        model.eval()
+        pbar = tqdm.tqdm(total=self.num_samples // self.batch_size + 1)
+        with torch.no_grad():
+            while generated_count < self.num_samples:
+
+                next_batch_size = (self.num_samples - generated_count) % self.batch_size
+                if next_batch_size == 0:
+                    next_batch_size = self.batch_size
+
+
+                z = torch.randn((next_batch_size, self.latent_dim))
+                images, labels = model.generate(z)
+
+                cl_labels = self.classifier(images)
+                cl_labels = torch.argmax(cl_labels, dim=1)
+
+                ce_loss = self.ce(labels, cl_labels)
+
+
+                reconstruction = self.ae(images)
+                recon_losses = self.bce(reconstruction, images)
+
+                loss = self.bce_weight * recon_losses.item() + self.ce_weight * ce_loss.item()
+
+                acc += loss
+
+                generated_count += next_batch_size
+
+                pbar.update(1)
+        self.result = acc
+
+    def evaluate_batch(self, model, batch):
+        return 0
+
+    def accumulate_result(self, result):
+        return
+
+    def get_result(self):
+        return self.result
+
+def kl_annealing(epoch, epochs, model, **kwargs):
+    model.beta = model.beta + (kwargs["final_beta"] - kwargs["start_beta"]) / (epochs - kwargs["beta_start_delay"])
+
+
+if __name__ == '__main__':
+
+    hparams = {
+        "learning_rate": ((1e-4, 1e-2), "log"),
+        "batch_size": 64,
+        "epochs": 10,
+        "patience": 5,
+        "optimizer_method": "Adam",
+        # "optimizer_kwargs": {"weight_decay": 1e-4},
+        "loss_func": nn.BCELoss(reduction="sum"),
+        "alpha": ((0, 10), "float"),
+        "beta": ((0, 10), "float"),
+        "classifier_loss_weight": ((0, 10), "float"),
+        "start_beta": ((0, 1), "float"),
+        "final_beta": ((1, 10), "float"),
+        "beta_start_delay": ((4, 12), "int"),
+        "feature_space": ([(64, 7, 7)], "item"),
+        "latent_dim": 2,
+        "sample_input_shape": ([(1, 1, 28, 28)], "item"),
+        "classifier_num_layers": 3,
+        "use_scheduler": ((True, False, False), "item"),
+        "use_grad_clip": ((True, False, False), "item"),
+        "gamma": ((0.8, 0.95), "log"),
+        "epoch_function": ((kl_annealing, None, None), "item"),
+    }
+
+
+    train_loader, val_loader, test_loader = get_mnist_loaders(hparams["batch_size"])
+
+    evaluator = GenerationEvaluator()
+
+    tuner = Tuner(ImageVAESemiSupervised)
+
+    num_search = 500
+    ranks_considered = 15
+    check_multiplier = 3
+
+    # train model
+    model, best_config, results = tuner.tune(train_loader, val_loader, hparams,
+                                             evaluator=evaluator,
+                                             num_search=num_search,
+                                             ranks_considered=ranks_considered,
+                                             check_multiplicant=check_multiplier,
+                                             mode='dynamic', model_path='models/mnist_vae_semi_tuned_')
+
+    model.eval()
+
+    model.save_model("models/mnist_vae_semi_tuned_best.pt", best_config)
+
+    plot_reconstructions(model, test_loader)
+
+    distribution = plot_label_clusters(model.encoder, test_loader)
+
+    plot_latent(model.decoder, dist=distribution, stds=2)