82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
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')
|