Done: First VIsualization

ToDo: Visualization for all classes, latent space setups

networks/adverserial_auto_encoder.py

@@ -25,6 +25,10 @@ class AdversarialAutoEncoder(AutoEncoder):
 
 class AdversarialAELightningOverrides:
 
+    @property
+    def name(self):
+        return self.__class__.__name__
+
     def forward(self, x):
         return self.network.forward(x)
 
@@ -46,6 +50,7 @@ class AdversarialAELightningOverrides:
             d_loss_fake = mse_loss(d_fake_prediction, torch.ones(d_fake_prediction.shape))
 
             # Calculate the mean over both the real and the fake acc
+            # ToDo: do i need to compute this seperate?
             d_loss = 0.5 * torch.add(d_loss_real, d_loss_fake)
             return {'loss': d_loss}
 

networks/auto_encoder.py

@@ -5,16 +5,12 @@ from torch import Tensor
 
 #######################
 # Basic AE-Implementation
-class AutoEncoder(Module, ABC):
+class AutoEncoder(AbstractNeuralNetwork, ABC):
 
-    @property
-    def name(self):
-        return self.__class__.__name__
-
-    def __init__(self, dataParams, **kwargs):
+    def __init__(self, latent_dim: int, dataParams: dict, **kwargs):
         super(AutoEncoder, self).__init__()
         self.dataParams = dataParams
-        self.latent_dim = kwargs.get('latent_dim', 2)
+        self.latent_dim = latent_dim
         self.encoder = Encoder(self.latent_dim)
         self.decoder = Decoder(self.latent_dim, self.dataParams['features'])
 
@@ -31,6 +27,10 @@ class AutoEncoder(Module, ABC):
 
 class AutoEncoderLightningOverrides:
 
+    @property
+    def name(self):
+        return self.__class__.__name__
+
     def forward(self, x):
         return self.network.forward(x)
 

networks/modules.py

@@ -1,9 +1,24 @@
 import torch
 import pytorch_lightning as pl
-from torch.nn import Module, Linear, ReLU, Tanh, Sigmoid, Dropout, GRU
+from torch.nn import Module, Linear, ReLU, Tanh, Sigmoid, Dropout, GRU, AvgPool2d
 
 from abc import ABC, abstractmethod
 
+#######################
+# Abstract NN Class
+
+class AbstractNeuralNetwork(Module):
+
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    def __init__(self):
+        super(AbstractNeuralNetwork, self).__init__()
+
+    def forward(self, batch):
+        pass
+
 
 ######################
 # Abstract Network class following the Lightning Syntax
@@ -102,6 +117,15 @@ class RNNOutputFilter(Module):
         return out if not self.only_last else out[:, -1, :]
 
 
+class AvgDimPool(Module):
+
+    def __init__(self):
+        super(AvgDimPool, self).__init__()
+
+    def forward(self, x):
+        return x.mean(-2)
+
+
 #######################
 # Network Modules
 # Generators, Decoders, Encoders, Discriminators
@@ -112,8 +136,8 @@ class Discriminator(Module):
         self.dataParams = dataParams
         self.latent_dim = latent_dim
         self.l1 = Linear(self.latent_dim, self.dataParams['features'] * 10)
-        self.l2 = Linear(self.dataParams['features']*10, self.dataParams['features'] * 20)
-        self.lout = Linear(self.dataParams['features']*20, 1)
+        self.l2 = Linear(self.dataParams['features'] * 10, self.dataParams['features'] * 20)
+        self.lout = Linear(self.dataParams['features'] * 20, 1)
         self.dropout = Dropout(dropout)
         self.activation = activation()
         self.sigmoid = Sigmoid()
@@ -149,6 +173,7 @@ class EncoderLinearStack(Module):
 
     def __init__(self):
         super(EncoderLinearStack, self).__init__()
+        # FixMe: Get Hardcoded shit out of here
         self.l1 = Linear(6, 100, bias=True)
         self.l2 = Linear(100, 10, bias=True)
         self.activation = ReLU()
@@ -188,6 +213,31 @@ class Encoder(Module):
         return tensor
 
 
+class PoolingEncoder(Module):
+
+    def __init__(self, lat_dim, variational=False):
+        self.lat_dim = lat_dim
+        self.variational = variational
+
+        super(PoolingEncoder, self).__init__()
+        self.p = AvgDimPool()
+        self.l = EncoderLinearStack()
+        if variational:
+            self.mu = Linear(10, self.lat_dim)
+            self.logvar = Linear(10, self.lat_dim)
+        else:
+            self.lat_dim_layer = Linear(10, self.lat_dim)
+
+    def forward(self, x):
+        tensor = self.p(x)
+        tensor = self.l(tensor)
+        if self.variational:
+            tensor = self.mu(tensor), self.logvar(tensor)
+        else:
+            tensor = self.lat_dim_layer(tensor)
+        return tensor
+
+
 class Decoder(Module):
 
     def __init__(self, latent_dim, *args, variational=False):

networks/seperating_adversarial_auto_encoder.py

@@ -0,0 +1,96 @@
+from networks.auto_encoder import AutoEncoder
+from torch.nn.functional import mse_loss
+from networks.modules import *
+import torch
+
+
+class SeperatingAdversarialAutoEncoder(Module):
+
+    def __init__(self, latent_dim, dataParams, **kwargs):
+        assert latent_dim % 2 == 0, f'Your latent space needs to be even, not odd, but was: "{latent_dim}"'
+        super(SeperatingAdversarialAutoEncoder, self).__init__()
+
+        self.latent_dim = latent_dim
+        self.dataParams = dataParams
+        self.spatial_encoder = PoolingEncoder(self.latent_dim // 2)
+        self.temporal_encoder = Encoder(self.latent_dim // 2)
+        self.decoder = Decoder(self.latent_dim, self.dataParams['features'])
+        self.spatial_discriminator = Discriminator(self.latent_dim // 2, self.dataParams)
+        self.temporal_discriminator = Discriminator(self.latent_dim // 2, self.dataParams)
+
+    def forward(self, batch):
+        # Encoder
+        #  outputs, hidden (Batch, Timesteps aka. Size, Features / Latent Dim Size)
+        z_spatial, z_temporal = self.spatial_encoder(batch), self.temporal_encoder(batch)
+        # Decoder
+        # First repeat the data accordingly to the batch size
+        z_concat = torch.cat((z_spatial, z_temporal), dim=-1)
+        z_repeatet = Repeater((batch.shape[0], self.dataParams['size'], -1))(z_concat)
+        x_hat = self.decoder(z_repeatet)
+        return z_spatial, z_temporal, x_hat
+
+
+class SeparatingAdversarialAELightningOverrides:
+
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    def forward(self, x):
+        return self.network.forward(x)
+
+    def training_step(self, batch, _, optimizer_i):
+        spatial_latent_fake, temporal_latent_fake, batch_hat = self.network.forward(batch)
+        if optimizer_i == 0:
+            # ---------------------
+            #  Train temporal Discriminator
+            # ---------------------
+            # latent_fake, reconstruction
+            temporal_latent_real = self.normal.sample(temporal_latent_fake.shape)
+
+            # Evaluate the input
+            temporal_real_prediction = self.network.temporal_discriminator.forward(temporal_latent_real)
+            temporal_fake_prediction = self.network.temporal_discriminator.forward(temporal_latent_fake)
+
+            # Train the discriminator
+            temporal_loss_real = mse_loss(temporal_real_prediction, torch.zeros(temporal_real_prediction.shape))
+            temporal_loss_fake = mse_loss(temporal_fake_prediction, torch.ones(temporal_fake_prediction.shape))
+
+            # Calculate the mean over bot the real and the fake acc
+            # ToDo: do i need to compute this seperate?
+            d_loss = 0.5 * torch.add(temporal_loss_real, temporal_loss_fake)
+            return {'loss': d_loss}
+
+        if optimizer_i == 1:
+            # ---------------------
+            #  Train spatial Discriminator
+            # ---------------------
+            # latent_fake, reconstruction
+            spatial_latent_real = self.normal.sample(spatial_latent_fake.shape)
+
+            # Evaluate the input
+            spatial_real_prediction = self.network.spatial_discriminator.forward(spatial_latent_real)
+            spatial_fake_prediction = self.network.spatial_discriminator.forward(spatial_latent_fake)
+
+            # Train the discriminator
+            spatial_loss_real = mse_loss(spatial_real_prediction, torch.zeros(spatial_real_prediction.shape))
+            spatial_loss_fake = mse_loss(spatial_fake_prediction, torch.ones(spatial_fake_prediction.shape))
+
+            # Calculate the mean over bot the real and the fake acc
+            # ToDo: do i need to compute this seperate?
+            d_loss = 0.5 * torch.add(spatial_loss_real, spatial_loss_fake)
+            return {'loss': d_loss}
+
+        elif optimizer_i == 2:
+            # ---------------------
+            #  Train AutoEncoder
+            # ---------------------
+            loss = mse_loss(batch, batch_hat)
+            return {'loss': loss}
+
+        else:
+            raise RuntimeError('This should not have happened, catch me if u can.')
+
+
+if __name__ == '__main__':
+    raise PermissionError('Get out of here - never run this module')

networks/variational_auto_encoder.py

@@ -4,7 +4,7 @@ from torch.nn.functional import mse_loss
 
 #######################
 # Basic AE-Implementation
-class VariationalAutoEncoder(Module, ABC):
+class VariationalAutoEncoder(AbstractNeuralNetwork, ABC):
 
     @property
     def name(self):
@@ -34,6 +34,10 @@ class VariationalAutoEncoder(Module, ABC):
 
 class VariationalAutoEncoderLightningOverrides:
 
+    @property
+    def name(self):
+        return self.network.name
+
     def forward(self, x):
         return self.network.forward(x)