Merge pull request #16 from renan-siqueira/develop

Merge develop into main

Author: renan-siqueira

json/params.json

@@ -2,8 +2,8 @@
     "batch_size": 32,
     "resolution": 64,
     "encoding_dim": 16,
-    "num_epochs": 1000,
+    "num_epochs": 25,
     "learning_rate": 0.001,
-    "ae_type": "conv",
+    "ae_type": "conv_dae",
     "save_checkpoint": null
 }

models/__init__.py

@@ -1,4 +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 .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/autoencoder_dae.py

@@ -0,0 +1,48 @@
+import torch.nn as nn
+import torch
+
+
+class DenoisingAutoencoder(nn.Module):
+    def __init__(self, input_dim, encoding_dim):
+        super(DenoisingAutoencoder, self).__init__()
+
+        self.model_structure = 'linear'
+        self.model_variant = 'vanilla'
+
+        self.encoder = nn.Sequential(
+            nn.Linear(input_dim, 1024),
+            nn.ReLU(),
+            nn.Linear(1024, 512),
+            nn.ReLU(),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.Linear(128, 64),
+            nn.ReLU(),
+            nn.Linear(64, encoding_dim)
+        )
+
+        self.decoder = nn.Sequential(
+            nn.Linear(encoding_dim, 64),
+            nn.ReLU(),
+            nn.Linear(64, 128),
+            nn.ReLU(),
+            nn.Linear(128, 256),
+            nn.ReLU(),
+            nn.Linear(256, 512),
+            nn.ReLU(),
+            nn.Linear(512, 1024),
+            nn.ReLU(),
+            nn.Linear(1024, input_dim),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        noise = torch.randn_like(x) * 0.1
+        x_corrupted = x + noise
+
+        x_encoded = self.encoder(x_corrupted)
+        x_decoded = self.decoder(x_encoded)
+
+        return x_decoded

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,18 +2,47 @@
 import json
 import random
 import time
+import argparse
 
 import numpy as np
 import torch
 
-from models import Autoencoder, ConvolutionalAutoencoder, ConvolutionalVAE, VariationalAutoencoder
+from models import (
+    Autoencoder,
+    ConvolutionalAutoencoder,
+    ConvolutionalVAE,
+    DenoisingAutoencoder,
+    VariationalAutoencoder,
+    DenoisingConvolutionalAutoencoder,
+)
+
 from settings import settings
 from utils.dataloader import get_dataloader
 from utils.trainer import train_autoencoder, visualize_reconstructions, load_checkpoint, evaluate_autoencoder
 from utils import utils
 
 
-def set_seed(seed=42):
+def get_model_by_type(ae_type=None, input_dim=None, encoding_dim=None, device=None):
+    models = {
+        'ae': lambda: Autoencoder(input_dim, encoding_dim),
+        '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}")
+
+    model = models[ae_type]()
+    return model.to(device)
+
+
+def set_seed(seed):
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     np.random.seed(seed)
@@ -28,35 +57,28 @@ 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
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     dataloader = get_dataloader(settings.DATA_PATH, batch_size, resolution)
 
-    if ae_type == 'ae':
-        model = Autoencoder(input_dim, encoding_dim).to(device)
-    elif ae_type == 'conv':
-        model = ConvolutionalAutoencoder().to(device)
-    elif ae_type == 'vae':
-        model = VariationalAutoencoder().to(device)
-    elif ae_type == 'conv_vae':
-        model = ConvolutionalVAE().to(device)
-    else:
-        raise ValueError(f"Unknown AE type: {ae_type}")
-
+    model = get_model_by_type(ae_type, input_dim, encoding_dim, device)
     optimizer = torch.optim.Adam(model.parameters())
 
     start_epoch = 0
@@ -78,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)