67 lines
2.3 KiB
Python
67 lines
2.3 KiB
Python
from networks.auto_encoder import AutoEncoder
|
|
from torch.nn.functional import mse_loss
|
|
from torch.nn import Sequential, Linear, ReLU, Dropout, Sigmoid
|
|
from torch.distributions import Normal
|
|
from networks.modules import *
|
|
import torch
|
|
|
|
|
|
class AdversarialAutoEncoder(AutoEncoder):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super(AdversarialAutoEncoder, self).__init__(*args, **kwargs)
|
|
self.discriminator = Discriminator(self.latent_dim, self.dataParams)
|
|
|
|
def forward(self, batch):
|
|
# Encoder
|
|
# outputs, hidden (Batch, Timesteps aka. Size, Features / Latent Dim Size)
|
|
z = self.encoder(batch)
|
|
# Decoder
|
|
# First repeat the data accordingly to the batch size
|
|
z_repeatet = Repeater((batch.shape[0], self.dataParams['size'], -1))(z)
|
|
x_hat = self.decoder(z_repeatet)
|
|
return z, x_hat
|
|
|
|
|
|
class AdversarialAELightningOverrides:
|
|
|
|
def forward(self, x):
|
|
return self.network.forward(x)
|
|
|
|
def training_step(self, batch, _, optimizer_i):
|
|
if optimizer_i == 0:
|
|
# ---------------------
|
|
# Train Discriminator
|
|
# ---------------------
|
|
# latent_fake, reconstruction
|
|
latent_fake, _ = self.network.forward(batch)
|
|
latent_real = self.normal.sample(latent_fake.shape)
|
|
|
|
# Evaluate the input
|
|
d_real_prediction = self.network.discriminator.forward(latent_real)
|
|
d_fake_prediction = self.network.discriminator.forward(latent_fake)
|
|
|
|
# Train the discriminator
|
|
d_loss_real = mse_loss(d_real_prediction, torch.zeros(d_real_prediction.shape))
|
|
d_loss_fake = mse_loss(d_fake_prediction, torch.ones(d_fake_prediction.shape))
|
|
|
|
# Calculate the mean over both the real and the fake acc
|
|
d_loss = 0.5 * torch.add(d_loss_real, d_loss_fake)
|
|
return {'loss': d_loss}
|
|
|
|
elif optimizer_i == 1:
|
|
# ---------------------
|
|
# Train AutoEncoder
|
|
# ---------------------
|
|
# z, x_hat
|
|
_, batch_hat = self.forward(batch)
|
|
loss = mse_loss(batch, batch_hat)
|
|
return {'loss': loss}
|
|
|
|
else:
|
|
raise RuntimeError('This should not have happened, catch me if u can.')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
raise PermissionError('Get out of here - never run this module')
|