fix noisy encoder issue

autoencoders/conf/comps.yaml

@@ -46,7 +46,7 @@ models:
         - 28
       loss_func:
         _target_: torch.nn.MSELoss
-    ckpt_path: DAE/train/2023-08-29/21-19-54/checkpoints/epoch=10-step=2585.ckpt
+    ckpt_path: DAE/train/2023-09-09/14-08-48/checkpoints/epoch=11-step=2820.ckpt
   deep-ae:
     name: DeepAutoEncoder
     module:
@@ -60,16 +60,13 @@ models:
   deep-dae:
     name: DeepDenoisingAutoEncoder
     module:
-      _target_: autoencoders.models.deep_ae.DeepAutoEncoder
+      _target_: autoencoders.models.deep_ae.DeepDenoisingAutoEncoder
       base_channels: 16
       latent_dim: 256
       input_channels: 1
       loss_func:
         _target_: torch.nn.MSELoss
-      encoder:
-        _target_: autoencoders.modules.NoisyCNNEncoder
-        _partial_: true
-    ckpt_path: deep-DAE/train/2023-09-07/12-56-27/checkpoints/epoch=5-step=1410.ckpt
+    ckpt_path: deep-DAE/train/2023-09-09/14-17-53/checkpoints/epoch=5-step=1410.ckpt
   vae:
     name: VariationalAutoEncoder
     module:

autoencoders/conf/model/deep_dae.yaml

@@ -9,12 +9,9 @@ defaults:
 name: deep-DAE
 
 nn:
-  _target_: autoencoders.models.deep_ae.DeepAutoEncoder
+  _target_: autoencoders.models.deep_ae.DeepDenoisingAutoEncoder
   base_channels: 16
   latent_dim: 256
   input_channels: 1
   loss_func:
     _target_: torch.nn.MSELoss
-  encoder:
-    _target_: autoencoders.modules.NoisyCNNEncoder
-    _partial_: True

autoencoders/models/base_dae.py

@@ -19,8 +19,20 @@ def __init__(
         super().__init__(layers, input_shape, loss_func, optim, scheduler)
         self.encoder = nn.Sequential(
             nn.Flatten(),
-            WhiteNoise(),
             nn.Linear(self.input_shape[0] * self.input_shape[1], self.encoder_shape[0]),
             *self.encoder_layers,
             nn.Linear(self.encoder_shape[-2], self.encoder_shape[-1]),
         )
+
+        self.augment = WhiteNoise()
+
+    def training_step(self, batch, idx):
+        original = batch[0]
+        augmented = self.augment(original)
+        reconstructed = self(augmented)
+        loss = self.loss_func(original, reconstructed)
+
+        self.log("loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=False)
+        metrics = {"train-loss": loss}
+        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=False, logger=True)
+        return loss

autoencoders/models/deep_ae.py

@@ -4,7 +4,7 @@
 import torch.nn as nn
 
 from autoencoders.models.base import BaseModule
-from autoencoders.modules import CNNDecoder, CNNEncoder
+from autoencoders.modules import CNNDecoder, CNNEncoder, WhiteNoise
 
 
 class DeepAutoEncoder(BaseModule):
@@ -41,3 +41,33 @@ def training_step(self, batch, idx):
         metrics = {"train-loss": loss}
         self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=False, logger=True)
         return loss
+
+
+class DeepDenoisingAutoEncoder(DeepAutoEncoder):
+    def __init__(
+        self,
+        base_channels: int,
+        latent_dim: int,
+        encoder: nn.Module = CNNEncoder,
+        decoder: nn.Module = CNNDecoder,
+        input_channels: int = 1,
+        loss_func: nn.Module = nn.MSELoss(),
+        optim: Optional[torch.optim.Optimizer] = None,
+        scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
+    ):
+        super().__init__(
+            base_channels, latent_dim, encoder, decoder, input_channels, loss_func, optim, scheduler
+        )
+
+        self.augment = WhiteNoise()
+
+    def training_step(self, batch, idx):
+        original = batch[1]
+        augmented = self.augment(batch[0])
+        reconstructed = self(augmented)
+        loss = self.loss_func(original, reconstructed)
+
+        self.log("loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=False)
+        metrics = {"train-loss": loss}
+        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=False, logger=True)
+        return loss

outputs/DAE/train/2023-08-29/21-19-54/.hydra/hydra.yaml → outputs/DAE/train/2023-09-09/14-08-48/.hydra/hydra.yaml

@@ -152,7 +152,7 @@ hydra:
     - path: ''
       schema: structured
       provider: schema
-    output_dir: /Users/chrissantiago/Dropbox/GitHub/autoencoders/outputs/DAE/train/2023-08-29/21-19-54
+    output_dir: /Users/chrissantiago/Dropbox/GitHub/autoencoders/outputs/DAE/train/2023-09-09/14-08-48
     choices:
       experiment: null
       callbacks: encoder

outputs/DAE/train/2023-09-09/14-08-48/checkpoints/best_k_models.yaml

@@ -0,0 +1 @@
+/Users/chrissantiago/Dropbox/GitHub/autoencoders/outputs/DAE/train/2023-09-09/14-08-48/checkpoints/epoch=11-step=2820.ckpt: 0.03247947618365288

outputs/comps.json

@@ -61,20 +61,20 @@
       8000
     ],
     "acc": [
-      0.3071,
-      0.6082,
-      0.6975,
-      0.7236,
-      0.7595,
-      0.7575
+      0.469,
+      0.8038,
+      0.8869,
+      0.9068,
+      0.9175,
+      0.9282
     ],
     "auc": [
-      0.6153,
-      0.7826,
-      0.8321,
-      0.8466,
-      0.8665,
-      0.8655
+      0.7053,
+      0.8912,
+      0.9372,
+      0.9483,
+      0.9543,
+      0.9602
     ]
   },
   "DeepAutoEncoder": {
@@ -113,20 +113,20 @@
       8000
     ],
     "acc": [
-      0.3758,
-      0.6647,
-      0.7636,
-      0.7967,
-      0.818,
-      0.8349
+      0.4723,
+      0.858,
+      0.9346,
+      0.9458,
+      0.951,
+      0.9613
     ],
     "auc": [
-      0.6535,
-      0.814,
-      0.8688,
-      0.8872,
-      0.899,
-      0.9084
+      0.7071,
+      0.9212,
+      0.9637,
+      0.9699,
+      0.9728,
+      0.9785
     ]
   },
   "VariationalAutoEncoder": {