feat: create denoising autoencoder model and refactory run.py to create smart function of model selection

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": "ae",
     "save_checkpoint": null
 }

models/__init__.py

@@ -1,4 +1,5 @@
 from .autoencoder import Autoencoder
 from .convolutional_autoencoder import ConvolutionalAutoencoder
 from .convolutional_vae import ConvolutionalVAE
+from .denoising_autoencoder import DenoisingAutoencoder
 from .variational_autoencoder import VariationalAutoencoder

models/denoising_autoencoder.py

@@ -0,0 +1,46 @@
+import torch.nn as nn
+import torch
+
+
+class DenoisingAutoencoder(nn.Module):
+    def __init__(self, input_dim, encoding_dim):
+        print('***** Denoising Autoencoder input_dim:', input_dim)
+        super(DenoisingAutoencoder, self).__init__()
+
+        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

run.py

@@ -6,14 +6,37 @@
 import numpy as np
 import torch
 
-from models import Autoencoder, ConvolutionalAutoencoder, ConvolutionalVAE, VariationalAutoencoder
+from models import (
+    Autoencoder,
+    ConvolutionalAutoencoder,
+    ConvolutionalVAE,
+    DenoisingAutoencoder,
+    VariationalAutoencoder,
+)
+
 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, input_dim, encoding_dim, device):
+    models = {
+        'ae': lambda: Autoencoder(input_dim, encoding_dim),
+        'conv': ConvolutionalAutoencoder,
+        'conv_vae': ConvolutionalVAE,
+        'dae': lambda: DenoisingAutoencoder(input_dim, encoding_dim),
+        'vae': VariationalAutoencoder,
+    }
+
+    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)
@@ -46,17 +69,7 @@ def main(load_trained_model):
     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