Lightning integration basic ae, dataloaders and dataset

networks/basic_ae.py

@@ -1,52 +1,73 @@
-from torch.nn import Sequential, Linear, GRU
-from data.dataset import DataContainer
-
+from torch.nn import Sequential, Linear, GRU, ReLU, Tanh
 from .modules import *
+from torch.nn.functional import mse_loss
+
 
 
 #######################
 # Basic AE-Implementation
 class BasicAE(Module, ABC):
 
+    @property
+    def name(self):
+        return self.__class__.__name__
+
     def __init__(self, dataParams, **kwargs):
         super(BasicAE, self).__init__()
         self.dataParams = dataParams
         self.latent_dim = kwargs.get('latent_dim', 2)
         self.encoder = self._build_encoder()
-        self.decoder = self._build_decoder()
-
+        self.decoder = self._build_decoder(out_shape=self.dataParams['features'])
 
     def _build_encoder(self):
-        encoder = Sequential()
-        encoder.add_module(f'EncoderLinear_{1}', Linear(6, 10, bias=True))
-        encoder.add_module(f'EncoderLinear_{2}', Linear(10, 10, bias=True))
-        gru = Sequential()
-        gru.add_module('Encoder', TimeDistributed(encoder))
-        gru.add_module('GRU', GRU(10, self.latent_dim))
+        encoder = Sequential(
+            Linear(6, 100, bias=True),
+            ReLU(),
+            Linear(100, 10, bias=True),
+            ReLU()
+        )
+        gru = Sequential(
+            TimeDistributed(encoder),
+            GRU(10, 10, batch_first=True),
+            RNNOutputFilter(only_last=True),
+            Linear(10, self.latent_dim)
+        )
         return gru
 
-    def _build_decoder(self):
-        decoder = Sequential()
-        decoder.add_module(f'DecoderLinear_{1}', Linear(10, 10, bias=True))
-        decoder.add_module(f'DecoderLinear_{2}', Linear(10, self.dataParams['features'], bias=True))
+    def _build_decoder(self, out_shape):
+        decoder = Sequential(
+            Linear(10, 100, bias=True),
+            ReLU(),
+            Linear(100, out_shape, bias=True),
+            Tanh()
+        )
 
-        gru = Sequential()
-        # There needs to be ab propper bat
-        gru.add_module('Repeater', Repeater((1, self.dataParams['size'], -1)))
-        gru.add_module('GRU', GRU(self.latent_dim, 10))
-        gru.add_module('GRU Filter', RNNOutputFilter())
-        gru.add_module('Decoder', TimeDistributed(decoder))
+        gru = Sequential(
+            GRU(self.latent_dim, 10,batch_first=True),
+            RNNOutputFilter(),
+            TimeDistributed(decoder)
+        )
         return gru
 
-    def forward(self, batch):
-        batch_size = batch.shape[0]
-        self.decoder.Repeater.shape = (batch_size, ) + self.decoder.Repeater.shape[-2:]
+    def forward(self, batch: torch.Tensor):
+        # Encoder
         #  outputs, hidden (Batch, Timesteps aka. Size, Features / Latent Dim Size)
-        outputs, _ = self.encoder(batch)
-        z = outputs[:, -1]
+        z = self.encoder(batch)
+        # Decoder
+        # First repeat the data accordingly to the batch size
+        z = Repeater((batch.shape[0], self.dataParams['size'], -1))(z)
         x_hat = self.decoder(z)
         return z, x_hat
 
 
+class AELightningOverrides:
+
+    def training_step(self, x, batch_nb):
+        # z, x_hat
+        _, x_hat = self.forward(x)
+        loss = mse_loss(x, x_hat)
+        return {'loss': loss}
+
+
 if __name__ == '__main__':
     raise PermissionError('Get out of here - never run this module')

networks/basic_vae.py

@@ -0,0 +1,81 @@
+from torch.nn import Sequential, Linear, GRU, ReLU
+from .modules import *
+from torch.nn.functional import mse_loss
+
+
+#######################
+# Basic AE-Implementation
+class BasicVAE(Module, ABC):
+
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    def __init__(self, dataParams, **kwargs):
+        super(BasicVAE, self).__init__()
+        self.dataParams = dataParams
+        self.latent_dim = kwargs.get('latent_dim', 2)
+        self.encoder = self._build_encoder()
+        self.decoder = self._build_decoder(out_shape=self.dataParams['features'])
+        self.mu, self.logvar = Linear(10, self.latent_dim), Linear(10, self.latent_dim)
+
+    def _build_encoder(self):
+        linear_stack = Sequential(
+            Linear(6, 100, bias=True),
+            ReLU(),
+            Linear(100, 10, bias=True),
+            ReLU()
+        )
+        encoder = Sequential(
+            TimeDistributed(linear_stack),
+            GRU(10, 10, batch_first=True),
+            RNNOutputFilter(only_last=True),
+        )
+        return encoder
+
+    def reparameterize(self, mu, logvar):
+        # Lambda Layer, add gaussian noise
+        std = torch.exp(0.5*logvar)
+        eps = torch.randn_like(std)
+        return mu + eps*std
+
+    def _build_decoder(self, out_shape):
+        decoder = Sequential(
+            Linear(10, 100, bias=True),
+            ReLU(),
+            Linear(100, out_shape, bias=True),
+            ReLU()
+        )
+
+        sequential_decoder = Sequential(
+            GRU(self.latent_dim, 10, batch_first=True),
+            RNNOutputFilter(),
+            TimeDistributed(decoder)
+        )
+        return sequential_decoder
+
+    def forward(self, batch):
+        encoding = self.encoder(batch)
+        mu_logvar = self.mu(encoding), self.logvar(encoding)
+        z = self.reparameterize(*mu_logvar)
+        repeat = Repeater((batch.shape[0], self.dataParams['size'], -1))
+        x_hat = self.decoder(repeat(z))
+        return (x_hat, *mu_logvar)
+
+
+class VAELightningOverrides:
+
+    def training_step(self, x, batch_nb):
+        x_hat, logvar, mu = self.forward(x)
+        BCE = mse_loss(x_hat, x, reduction='mean')
+
+        # see Appendix B from VAE paper:
+        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
+        # https://arxiv.org/abs/1312.6114
+        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
+        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
+        return {'loss': BCE + KLD}
+
+
+if __name__ == '__main__':
+    raise PermissionError('Get out of here - never run this module')

networks/modules.py

@@ -90,13 +90,15 @@ class Repeater(Module):
 
 class RNNOutputFilter(Module):
 
-    def __init__(self, return_output=True):
+    def __init__(self, return_output=True, only_last=False):
         super(RNNOutputFilter, self).__init__()
+        self.only_last = only_last
         self.return_output = return_output
 
     def forward(self, x: tuple):
         outputs, hidden = x
-        return outputs if self.return_output else hidden
+        out = outputs if self.return_output else hidden
+        return out if not self.only_last else out[:, -1, :]
 
 
 if __name__ == '__main__':