Done: AE, VAE, AAE

ToDo: Double AAE, Visualization
All Modularized

networks/modules.py

@@ -1,6 +1,6 @@
 import torch
 import pytorch_lightning as pl
-from torch.nn import Module
+from torch.nn import Module, Linear, ReLU, Tanh, Sigmoid, Dropout, GRU
 
 from abc import ABC, abstractmethod
 
@@ -85,9 +85,10 @@ class Repeater(Module):
         self.shape = shape
 
     def forward(self, x: torch.Tensor):
-        x.unsqueeze_(-2)
+        x = x.unsqueeze(-2)
         return x.expand(self.shape)
 
+
 class RNNOutputFilter(Module):
 
     def __init__(self, return_output=True, only_last=False):
@@ -101,5 +102,108 @@ class RNNOutputFilter(Module):
         return out if not self.only_last else out[:, -1, :]
 
 
+#######################
+# Network Modules
+# Generators, Decoders, Encoders, Discriminators
+class Discriminator(Module):
+
+    def __init__(self, latent_dim, dataParams, dropout=.0, activation=ReLU):
+        super(Discriminator, self).__init__()
+        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.dropout = Dropout(dropout)
+        self.activation = activation()
+        self.sigmoid = Sigmoid()
+
+    def forward(self, x, **kwargs):
+        tensor = self.l1(x)
+        tensor = self.dropout(self.activation(tensor))
+        tensor = self.l2(tensor)
+        tensor = self.dropout(self.activation(tensor))
+        tensor = self.lout(tensor)
+        tensor = self.sigmoid(tensor)
+        return tensor
+
+
+class DecoderLinearStack(Module):
+
+    def __init__(self, out_shape):
+        super(DecoderLinearStack, self).__init__()
+        self.l1 = Linear(10, 100, bias=True)
+        self.l2 = Linear(100, out_shape, bias=True)
+        self.activation = ReLU()
+        self.activation_out = Tanh()
+
+    def forward(self, x):
+        tensor = self.l1(x)
+        tensor = self.activation(tensor)
+        tensor = self.l2(tensor)
+        tensor = self.activation_out(tensor)
+        return tensor
+
+
+class EncoderLinearStack(Module):
+
+    def __init__(self):
+        super(EncoderLinearStack, self).__init__()
+        self.l1 = Linear(6, 100, bias=True)
+        self.l2 = Linear(100, 10, bias=True)
+        self.activation = ReLU()
+
+    def forward(self, x):
+        tensor = self.l1(x)
+        tensor = self.activation(tensor)
+        tensor = self.l2(tensor)
+        tensor = self.activation(tensor)
+        return tensor
+
+
+class Encoder(Module):
+
+    def __init__(self, lat_dim, variational=False):
+        self.lat_dim = lat_dim
+        self.variational = variational
+
+        super(Encoder, self).__init__()
+        self.l_stack = TimeDistributed(EncoderLinearStack())
+        self.gru = GRU(10, 10, batch_first=True)
+        self.filter = RNNOutputFilter(only_last=True)
+        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.l_stack(x)
+        tensor = self.gru(tensor)
+        tensor = self.filter(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):
+        self.variational = variational
+        super(Decoder, self).__init__()
+        self.g = GRU(latent_dim, 10, batch_first=True)
+        self.filter = RNNOutputFilter()
+        self.l_stack = TimeDistributed(DecoderLinearStack(*args))
+        pass
+
+    def forward(self, x):
+        tensor = self.g(x)
+        tensor = self.filter(tensor)
+        tensor = self.l_stack(tensor)
+        return tensor
+
+
 if __name__ == '__main__':
     raise PermissionError('Get out of here - never run this module')