feat: NoiseLayer class + keep loss across trainings

src/easyvae/autoencoder.py

@@ -1,6 +1,6 @@
 import numpy as np
 from tqdm import tqdm
-from .layers import DeepNNLayer, SampleLayer
+from .layers import DeepNNLayer, SampleLayer, NoiseLayer
 from .activations import ActivationFunc, Identity
 from .plotters import Plotter, CAPlotter, VAEPlotter
 from .utils import interruptable
@@ -17,13 +17,15 @@ class AAutoencoder(ABC):
                  encoder_layers: list[int],
                  decoder_layers: list[int],
                  lr: float,
-                 activation_func: ActivationFunc):
+                 activation_func: ActivationFunc,
+                 noise=0):
         if encoder_layers[-1] != decoder_layers[0]:
             raise Exception(
                 f"Encoder output and decoder input don't match {encoder_layers[-1]} != {encoder_layers[0]}" # noqa
             )
         self.encoder = DeepNNLayer(encoder_layers, lr, activation_func)
         self.decoder = DeepNNLayer(decoder_layers, lr, activation_func)
+        self.noise = NoiseLayer(noise)
         self.space_dim = decoder_layers[0]
         self.lr = lr
         self.losses = [0]
@@ -78,8 +80,8 @@ class ClassicalAutoencoder(AAutoencoder):
         return loss / len(data_set)
 
     def train(self, v: np.ndarray):
-        out = self.decoder.forward(
-            self.encoder.forward(v)
+        out, _ = self.forward(
+            self.noise.forward(v)
         )
         error = out - v
         self.encoder.back(
@@ -96,6 +98,7 @@ class ClassicalAutoencoder(AAutoencoder):
                       patience: int,
                       display_loss: bool = False) -> list[float]:
         plotter = self.plotter_cls(self) if display_loss else Plotter(self)
+        if len(self.losses) == 0:
             self.losses = [self.loss(data_set)]
         epoch = 0
         no_improv = 0
@@ -167,7 +170,9 @@ class VariationalAutoencoder(AAutoencoder):
         return recon_loss, kl_loss
 
     def train(self, v: np.ndarray) -> tuple[float, float]:
-        out, _ = self.forward(v)
+        out, _ = self.forward(
+            self.noise.forward(v)
+        )
         error = out - v
         self.encoder.back(
             self.sampler.back(
@@ -186,6 +191,7 @@ class VariationalAutoencoder(AAutoencoder):
                       patience: int,
                       display_loss: bool = False) -> list[float]:
         plotter = self.plotter_cls(self) if display_loss else Plotter(self)
+        if len(self.recon_losses) == 0:
             recon_0, kl_0 = self.loss(data_set)
             self.recon_losses = [recon_0]
             self.KL_losses = [kl_0]

src/easyvae/layers.py

@@ -115,3 +115,12 @@ class DeepNNLayer:
         for layer in self.layers:
             layer.backprop()
 
+
+class NoiseLayer:
+    def __init__(self, amount=0.1):
+        self.amount = amount
+
+    def forward(self, v: np.ndarray):
+        if self.amount == 0:
+            return v
+        return v + np.random.normal(0, self.amount, v.shape)