feat: improvements on models and trainer. Create test mode to run with all models

Author: renan-siqueira

json/params.json

@@ -4,6 +4,6 @@
     "encoding_dim": 16,
     "num_epochs": 25,
     "learning_rate": 0.001,
-    "ae_type": "ae",
+    "ae_type": "conv_dae",
     "save_checkpoint": null
 }

models/__init__.py

@@ -1,5 +1,6 @@
 from .autoencoder import Autoencoder
+from .autoencoder_dae import DenoisingAutoencoder
+from .autoencoder_vae import VariationalAutoencoder
 from .convolutional_autoencoder import ConvolutionalAutoencoder
+from .convolutional_dae import DenoisingConvolutionalAutoencoder
 from .convolutional_vae import ConvolutionalVAE
-from .denoising_autoencoder import DenoisingAutoencoder
-from .variational_autoencoder import VariationalAutoencoder

models/autoencoder.py

@@ -3,8 +3,11 @@
 
 class Autoencoder(nn.Module):
     def __init__(self, input_dim, encoding_dim):
-        print('***** Autoencoder input_dim:', input_dim)
         super(Autoencoder, self).__init__()
+
+        self.model_structure = 'linear'
+        self.model_variant = 'vanilla'
+
         self.encoder = nn.Sequential(
             nn.Linear(input_dim, 1024),
             nn.ReLU(),
@@ -18,6 +21,7 @@ def __init__(self, input_dim, encoding_dim):
             nn.ReLU(),
             nn.Linear(64, encoding_dim)
         )
+
         self.decoder = nn.Sequential(
             nn.Linear(encoding_dim, 64),
             nn.ReLU(),

models/denoising_autoencoder.py models/autoencoder_dae.pymodels/denoising_autoencoder.py renamed to models/autoencoder_dae.py

@@ -4,9 +4,11 @@
 
 class DenoisingAutoencoder(nn.Module):
     def __init__(self, input_dim, encoding_dim):
-        print('***** Denoising Autoencoder input_dim:', input_dim)
         super(DenoisingAutoencoder, self).__init__()
 
+        self.model_structure = 'linear'
+        self.model_variant = 'vanilla'
+
         self.encoder = nn.Sequential(
             nn.Linear(input_dim, 1024),
             nn.ReLU(),

models/variational_autoencoder.py models/autoencoder_vae.pymodels/variational_autoencoder.py renamed to models/autoencoder_vae.py

@@ -7,6 +7,9 @@ class VariationalAutoencoder(nn.Module):
     def __init__(self, encoding_dim=128):
         super(VariationalAutoencoder, self).__init__()
 
+        self.model_structure = 'linear'
+        self.model_variant = 'vae'
+
         # Encoder
         self.enc1 = nn.Linear(3 * 64 * 64, 512)
         self.enc2 = nn.Linear(512, 256)

models/convolutional_autoencoder.py

@@ -7,6 +7,9 @@ class ConvolutionalAutoencoder(nn.Module):
     def __init__(self):
         super(ConvolutionalAutoencoder, self).__init__()
 
+        self.model_structure = 'convolutional'
+        self.model_variant = 'vanilla'
+
         # Encoder
         self.enc0 = nn.Conv2d(3, 256, kernel_size=3, padding=1)
         self.enc1 = nn.Conv2d(256, 128, kernel_size=3, padding=1)

models/convolutional_dae.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DenoisingConvolutionalAutoencoder(nn.Module):
+    def __init__(self):
+        super(DenoisingConvolutionalAutoencoder, self).__init__()
+
+        self.model_structure = 'convolutional'
+        self.model_variant = 'vanilla'
+        
+        # Encoder
+        self.enc0 = nn.Conv2d(3, 256, kernel_size=3, padding=1)
+        self.enc1 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
+        self.enc2 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+        self.enc3 = nn.Conv2d(64, 32, kernel_size=3, padding=1)
+        self.enc4 = nn.Conv2d(32, 16, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2, 2, return_indices=True)
+
+        # Decoder
+        self.dec0 = nn.ConvTranspose2d(16, 32, kernel_size=2, stride=2)
+        self.dec1 = nn.ConvTranspose2d(32, 64, kernel_size=2, stride=2)
+        self.dec2 = nn.ConvTranspose2d(64, 128, kernel_size=2, stride=2)
+        self.dec3 = nn.ConvTranspose2d(128, 256, kernel_size=2, stride=2)
+        self.dec4 = nn.ConvTranspose2d(256, 3, kernel_size=2, stride=2)
+
+    def forward(self, x):
+        noise = torch.randn_like(x) * 0.1
+        x_corrupted = x + noise
+
+        # Encoder
+        x, _ = self.pool(F.relu(self.enc0(x_corrupted)))
+        x, _ = self.pool(F.relu(self.enc1(x)))
+        x, _ = self.pool(F.relu(self.enc2(x)))
+        x, _ = self.pool(F.relu(self.enc3(x)))
+        x, _ = self.pool(F.relu(self.enc4(x)))
+
+        # Decoder
+        x = F.relu(self.dec0(x))
+        x = F.relu(self.dec1(x))
+        x = F.relu(self.dec2(x))
+        x = F.relu(self.dec3(x))
+        x = torch.sigmoid(self.dec4(x))
+
+        return x

models/convolutional_vae.py

@@ -7,6 +7,9 @@ class ConvolutionalVAE(nn.Module):
     def __init__(self):
         super(ConvolutionalVAE, self).__init__()
 
+        self.model_structure = 'convolutional'
+        self.model_variant = 'vae'
+
         # Encoder
         self.enc1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
         self.enc2 = nn.Conv2d(64, 32, kernel_size=3, padding=1)

run.py

@@ -2,6 +2,7 @@
 import json
 import random
 import time
+import argparse
 
 import numpy as np
 import torch
@@ -12,6 +13,7 @@
     ConvolutionalVAE,
     DenoisingAutoencoder,
     VariationalAutoencoder,
+    DenoisingConvolutionalAutoencoder,
 )
 
 from settings import settings
@@ -20,14 +22,18 @@
 from utils import utils
 
 
-def get_model_by_type(ae_type, input_dim, encoding_dim, device):
+def get_model_by_type(ae_type=None, input_dim=None, encoding_dim=None, device=None):
     models = {
         'ae': lambda: Autoencoder(input_dim, encoding_dim),
-        'conv': ConvolutionalAutoencoder,
-        'conv_vae': ConvolutionalVAE,
         'dae': lambda: DenoisingAutoencoder(input_dim, encoding_dim),
         'vae': VariationalAutoencoder,
+        'conv': ConvolutionalAutoencoder,
+        'conv_dae': DenoisingConvolutionalAutoencoder,
+        'conv_vae': ConvolutionalVAE,
     }
+    
+    if ae_type is None:
+        return list(models.keys())
 
     if ae_type not in models:
         raise ValueError(f"Unknown AE type: {ae_type}")
@@ -51,18 +57,21 @@ def load_params(path):
     return params
 
 
-def main(load_trained_model):
+def main(load_trained_model, ae_type=None, num_epochs=5, test_mode=True):
     set_seed(1)
     params = load_params(settings.PATH_PARAMS_JSON)
 
     batch_size = params["batch_size"]
     resolution = params["resolution"]
     encoding_dim = params["encoding_dim"]
-    num_epochs = params["num_epochs"]
     learning_rate = params.get("learning_rate", 0.001)
-    ae_type = params["ae_type"]
     save_checkpoint = params["save_checkpoint"]
 
+    if not ae_type:
+        ae_type = params["ae_type"]
+        num_epochs = params["num_epochs"]
+        test_mode = False
+
     # Calculate input_dim based on resolution
     input_dim = 3 * resolution * resolution
 
@@ -91,25 +100,39 @@ def main(load_trained_model):
                 device=device,
                 start_epoch=start_epoch,
                 optimizer=optimizer,
-                ae_type=ae_type,
                 save_checkpoint=save_checkpoint
             )
 
             elapsed_time = utils.format_time(time.time() - start_time)
             print(f"\nTraining took {elapsed_time}")
             print(f"Training complete up to epoch {num_epochs}!")
-    
+
     except KeyboardInterrupt:
         print("\nTraining interrupted by user.")
 
-    valid_dataloader = get_dataloader(settings.VALID_DATA_PATH, batch_size, resolution)
-    avg_valid_loss = evaluate_autoencoder(model, valid_dataloader, device, ae_type)
-    print(f"\nAverage validation loss: {avg_valid_loss:.4f}\n")
+    if not test_mode:
+        valid_dataloader = get_dataloader(settings.VALID_DATA_PATH, batch_size, resolution)
+        avg_valid_loss = evaluate_autoencoder(model, valid_dataloader, device, ae_type)
+        print(f"\nAverage validation loss: {avg_valid_loss:.4f}\n")
 
-    visualize_reconstructions(
-        model, valid_dataloader, num_samples=10, device=device, ae_type=ae_type, resolution=resolution
-    )
+        visualize_reconstructions(
+            model, valid_dataloader, num_samples=10,
+            device=device, ae_type=ae_type, resolution=resolution
+        )
 
 
 if __name__ == "__main__":
-    main(False)
+    parser = argparse.ArgumentParser(description='Training and testing autoencoders.')
+    parser.add_argument(
+        '--test', action='store_true', help='Run the test routine for all autoencoders.'
+    )
+
+    args = parser.parse_args()
+
+    if args.test:
+        ae_types = get_model_by_type()
+        for ae_type in ae_types:
+            print(f"\n===== Training {ae_type} =====\n")
+            main(load_trained_model=False, ae_type=ae_type)
+    else:
+        main(load_trained_model=False)

utils/trainer.py

@@ -36,7 +36,7 @@ def load_checkpoint(model, optimizer, path, device):
 
 def train_autoencoder(
         model, dataloader, num_epochs, learning_rate, device,
-        start_epoch, optimizer, ae_type, save_checkpoint
+        start_epoch, optimizer, save_checkpoint
     ):
     criterion = nn.MSELoss()
     if optimizer is None:
@@ -46,10 +46,10 @@ def train_autoencoder(
         for data in dataloader:
             img = data.to(device)
 
-            if ae_type not in ['conv', 'conv_vae']:
+            if model.model_structure == 'linear':
                 img = img.view(img.size(0), -1)
 
-            if ae_type in ['vae', 'conv_vae']:
+            if model.model_variant == 'vae':
                 recon_x, mu, log_var = model(img)
                 loss = loss_function_vae(recon_x, img, mu, log_var)
             else:
@@ -81,7 +81,7 @@ def evaluate_autoencoder(model, dataloader, device, ae_type):
         for data in dataloader:
             img = data.to(device)
 
-            if ae_type not in ['conv', 'conv_vae']:
+            if model.model_structure == 'linear':
                 img = img.view(img.size(0), -1)
 
             if ae_type in ['vae', 'conv_vae']:
@@ -99,10 +99,10 @@ def visualize_reconstructions(model, dataloader, num_samples=10, device='cpu', s
     samples = next(iter(dataloader))
     samples = samples[:num_samples].to(device)
 
-    if ae_type not in ['conv', 'conv_vae']:
+    if model.model_structure == 'linear':
         samples = samples.view(samples.size(0), -1)
 
-    if ae_type in ['vae', 'conv_vae']:
+    if model.model_variant == 'vae':
         reconstructions, _, _ = model(samples)
     else:
         reconstructions = model(samples)