62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
from torch.optim import Adam
|
|
|
|
from .modules import *
|
|
from torch.nn.functional import mse_loss
|
|
|
|
|
|
#######################
|
|
# Basic AE-Implementation
|
|
class VariationalAE(AbstractNeuralNetwork, ABC):
|
|
|
|
@property
|
|
def name(self):
|
|
return self.__class__.__name__
|
|
|
|
def __init__(self, latent_dim=0, features=0, use_norm=True, **kwargs):
|
|
assert latent_dim and features
|
|
super(VariationalAE, self).__init__()
|
|
self.features = features
|
|
self.latent_dim = latent_dim
|
|
self.encoder = Encoder(self.latent_dim, variational=True, use_norm=use_norm)
|
|
self.decoder = Decoder(self.latent_dim, self.features, variational=True, use_norm=use_norm)
|
|
|
|
@staticmethod
|
|
def reparameterize(mu, logvar):
|
|
# Lambda Layer, add gaussian noise
|
|
std = torch.exp(0.5*logvar)
|
|
eps = torch.randn_like(std)
|
|
return mu + eps*std
|
|
|
|
def forward(self, batch):
|
|
mu, logvar = self.encoder(batch)
|
|
z = self.reparameterize(mu, logvar)
|
|
repeat = Repeater((batch.shape[0], batch.shape[1], -1))
|
|
x_hat = self.decoder(repeat(z))
|
|
return mu, logvar, x_hat
|
|
|
|
|
|
class VAE_LO(LightningModuleOverrides):
|
|
|
|
def __init__(self, train_on_predictions=False):
|
|
super(VAE_LO, self).__init__()
|
|
self.train_on_predictions=train_on_predictions
|
|
|
|
def training_step(self, batch, _):
|
|
x, y = batch
|
|
mu, logvar, x_hat = self.forward(x)
|
|
BCE = mse_loss(y, x_hat) if self.train_on_predictions else mse_loss(x, x_hat)
|
|
|
|
# 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}
|
|
|
|
def configure_optimizers(self):
|
|
return [Adam(self.parameters(), lr=0.004)]
|
|
|
|
|
|
if __name__ == '__main__':
|
|
raise PermissionError('Get out of here - never run this module')
|