Done: First VIsualization

ToDo: Visualization for all classes, latent space setups

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):