From 18305a9e7e69ba2a08b0e48cdd5b162157910a15 Mon Sep 17 00:00:00 2001 From: Si11ium Date: Sat, 24 Aug 2019 19:05:46 +0200 Subject: [PATCH] All models running. --- dataset.py | 4 +- networks/adverserial_auto_encoder.py | 13 +++--- networks/modules.py | 4 +- .../seperating_adversarial_auto_encoder.py | 25 ++++++----- run_models.py | 42 ++++++++----------- viz/viz_latent.py | 2 +- 6 files changed, 45 insertions(+), 45 deletions(-) diff --git a/dataset.py b/dataset.py index 6d160dc..ccf5a6c 100644 --- a/dataset.py +++ b/dataset.py @@ -47,8 +47,8 @@ class AbstractDataset(ConcatDataset, ABC): # maps = ['hotel', 'tum','gallery', 'queens', 'oet'] @property def maps(self): - return ['test', 'test2'] - # return ['hotel', 'tum','gallery', 'queens', 'oet'] + # return ['test', 'test2'] + return ['hotel', 'tum','gallery', 'queens', 'oet'] @property @abstractmethod diff --git a/networks/adverserial_auto_encoder.py b/networks/adverserial_auto_encoder.py index faaeee8..7cc55bd 100644 --- a/networks/adverserial_auto_encoder.py +++ b/networks/adverserial_auto_encoder.py @@ -6,6 +6,9 @@ from networks.modules import * import torch +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + class AdversarialAutoEncoder(AutoEncoder): def __init__(self, *args, **kwargs): @@ -32,18 +35,18 @@ class AdversarialAELightningOverrides(LightningModuleOverrides): if optimizer_i == 0: # --------------------- # Train Discriminator - # --------------------- + # ---------------------p # latent_fake, reconstruction - latent_fake, _ = self.network.forward(batch) - latent_real = self.normal.sample(latent_fake.shape) + latent_fake = self.network.encoder.forward(batch) + latent_real = self.normal.sample(latent_fake.shape).to(device) # 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)) + d_loss_real = mse_loss(d_real_prediction, torch.zeros(d_real_prediction.shape, device=device)) + d_loss_fake = mse_loss(d_fake_prediction, torch.ones(d_fake_prediction.shape, device=device)) # Calculate the mean over both the real and the fake acc # ToDo: do i need to compute this seperate? diff --git a/networks/modules.py b/networks/modules.py index 81a02bb..17ba250 100644 --- a/networks/modules.py +++ b/networks/modules.py @@ -25,9 +25,9 @@ class LightningModuleOverrides: @data_loader def tng_dataloader(self): - num_workers = os.cpu_count() // 2 + num_workers = 0 # os.cpu_count() // 2 return DataLoader(DataContainer('data', self.size, self.step), - shuffle=True, batch_size=100, num_workers=num_workers) + shuffle=True, batch_size=10000, num_workers=num_workers) class AbstractNeuralNetwork(Module): diff --git a/networks/seperating_adversarial_auto_encoder.py b/networks/seperating_adversarial_auto_encoder.py index b0872f4..556da2c 100644 --- a/networks/seperating_adversarial_auto_encoder.py +++ b/networks/seperating_adversarial_auto_encoder.py @@ -1,21 +1,22 @@ from torch.optim import Adam - -from networks.auto_encoder import AutoEncoder from torch.nn.functional import mse_loss from networks.modules import * import torch +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + class SeperatingAdversarialAutoEncoder(Module): - def __init__(self, latent_dim, features, **kwargs): + def __init__(self, latent_dim, features): super(SeperatingAdversarialAutoEncoder, self).__init__() self.latent_dim = latent_dim self.features = features self.spatial_encoder = PoolingEncoder(self.latent_dim) self.temporal_encoder = Encoder(self.latent_dim) - self.decoder = Decoder(self.latent_dim, self.features) + self.decoder = Decoder(self.latent_dim * 2, self.features) self.spatial_discriminator = Discriminator(self.latent_dim, self.features) self.temporal_discriminator = Discriminator(self.latent_dim, self.features) @@ -43,15 +44,17 @@ class SeparatingAdversarialAELightningOverrides(LightningModuleOverrides): # Train temporal Discriminator # --------------------- # latent_fake, reconstruction - temporal_latent_real = self.normal.sample(temporal_latent_fake.shape) + temporal_latent_real = self.normal.sample(temporal_latent_fake.shape).to(device) # Evaluate the input temporal_real_prediction = self.network.temporal_discriminator.forward(temporal_latent_real) temporal_fake_prediction = self.network.temporal_discriminator.forward(temporal_latent_fake) # Train the discriminator - temporal_loss_real = mse_loss(temporal_real_prediction, torch.zeros(temporal_real_prediction.shape)) - temporal_loss_fake = mse_loss(temporal_fake_prediction, torch.ones(temporal_fake_prediction.shape)) + temporal_loss_real = mse_loss(temporal_real_prediction, + torch.zeros(temporal_real_prediction.shape, device=device)) + temporal_loss_fake = mse_loss(temporal_fake_prediction, + torch.ones(temporal_fake_prediction.shape, device=device)) # Calculate the mean over bot the real and the fake acc # ToDo: do i need to compute this seperate? @@ -63,15 +66,17 @@ class SeparatingAdversarialAELightningOverrides(LightningModuleOverrides): # Train spatial Discriminator # --------------------- # latent_fake, reconstruction - spatial_latent_real = self.normal.sample(spatial_latent_fake.shape) + spatial_latent_real = self.normal.sample(spatial_latent_fake.shape).to(device) # Evaluate the input spatial_real_prediction = self.network.spatial_discriminator.forward(spatial_latent_real) spatial_fake_prediction = self.network.spatial_discriminator.forward(spatial_latent_fake) # Train the discriminator - spatial_loss_real = mse_loss(spatial_real_prediction, torch.zeros(spatial_real_prediction.shape)) - spatial_loss_fake = mse_loss(spatial_fake_prediction, torch.ones(spatial_fake_prediction.shape)) + spatial_loss_real = mse_loss(spatial_real_prediction, + torch.zeros(spatial_real_prediction.shape, device=device)) + spatial_loss_fake = mse_loss(spatial_fake_prediction, + torch.ones(spatial_fake_prediction.shape, device=device)) # Calculate the mean over bot the real and the fake acc # ToDo: do i need to compute this seperate? diff --git a/run_models.py b/run_models.py index 228430d..8fafd41 100644 --- a/run_models.py +++ b/run_models.py @@ -1,37 +1,33 @@ from torch.distributions import Normal - from networks.auto_encoder import * -import os + import time from networks.variational_auto_encoder import * from networks.adverserial_auto_encoder import * from networks.seperating_adversarial_auto_encoder import * from networks.modules import LightningModule -from torch.optim import Adam -from torch.utils.data import DataLoader -from pytorch_lightning import data_loader -from dataset import DataContainer - -from torch.nn import BatchNorm1d from pytorch_lightning import Trainer from test_tube import Experiment from argparse import Namespace from argparse import ArgumentParser +from distutils.util import strtobool args = ArgumentParser() -args.add_argument('--step', default=0) -args.add_argument('--features', default=0) -args.add_argument('--size', default=0) -args.add_argument('--latent_dim', default=0) +args.add_argument('--step', default=6) +args.add_argument('--features', default=6) +args.add_argument('--size', default=9) +args.add_argument('--latent_dim', default=4) args.add_argument('--model', default='Model') +args.add_argument('--refresh', type=strtobool, default=False) + # ToDo: How to implement this better? # other_classes = [AutoEncoder, AutoEncoderLightningOverrides] class Model(AutoEncoderLightningOverrides, LightningModule): - def __init__(self, parameters, **kwargs): + def __init__(self, parameters): assert all([x in parameters for x in ['step', 'size', 'latent_dim', 'features']]) self.size = parameters.size self.latent_dim = parameters.latent_dim @@ -43,7 +39,7 @@ class Model(AutoEncoderLightningOverrides, LightningModule): class AdversarialModel(AdversarialAELightningOverrides, LightningModule): - def __init__(self, parameters: Namespace, **kwargs): + def __init__(self, parameters: Namespace): assert all([x in parameters for x in ['step', 'size', 'latent_dim', 'features']]) self.size = parameters.size self.latent_dim = parameters.latent_dim @@ -57,7 +53,7 @@ class AdversarialModel(AdversarialAELightningOverrides, LightningModule): class SeparatingAdversarialModel(SeparatingAdversarialAELightningOverrides, LightningModule): - def __init__(self, parameters: Namespace, **kwargs): + def __init__(self, parameters: Namespace): assert all([x in parameters for x in ['step', 'size', 'latent_dim', 'features']]) self.size = parameters.size self.latent_dim = parameters.latent_dim @@ -65,16 +61,12 @@ class SeparatingAdversarialModel(SeparatingAdversarialAELightningOverrides, Ligh self.step = parameters.step super(SeparatingAdversarialModel, self).__init__() self.normal = Normal(0, 1) - self.network = SeperatingAdversarialAutoEncoder(self.latent_dim, self.features, **kwargs) + self.network = SeperatingAdversarialAutoEncoder(self.latent_dim, self.features) pass if __name__ == '__main__': - features = 6 - tag_dict = dict(features=features, latent_dim=4, size=5, step=6, refresh=False, - transforms=[BatchNorm1d(features)]) arguments = args.parse_args() - arguments.__dict__.update(tag_dict) model = globals()[arguments.model](arguments) @@ -82,19 +74,19 @@ if __name__ == '__main__': outpath = os.path.join(os.getcwd(), 'output', model.name, time.asctime().replace(' ', '_').replace(':', '-')) os.makedirs(outpath, exist_ok=True) exp = Experiment(save_dir=outpath) - exp.tag(tag_dict=tag_dict) + exp.tag(tag_dict=arguments.__dict__) from pytorch_lightning.callbacks import ModelCheckpoint checkpoint_callback = ModelCheckpoint( filepath=os.path.join(outpath, 'weights.ckpt'), - save_best_only=True, + save_best_only=False, verbose=True, - monitor='val_loss', # val_loss - mode='min', + period=4 ) - trainer = Trainer(experiment=exp, checkpoint_callback=checkpoint_callback, max_nb_epochs=15) # gpus=[0...LoL] + trainer = Trainer(experiment=exp, max_nb_epochs=250, gpus=[0], + add_log_row_interval=1000, checkpoint_callback=checkpoint_callback) trainer.fit(model) trainer.save_checkpoint(os.path.join(outpath, 'weights.ckpt')) diff --git a/viz/viz_latent.py b/viz/viz_latent.py index 15e61f0..84ba2a2 100644 --- a/viz/viz_latent.py +++ b/viz/viz_latent.py @@ -62,7 +62,7 @@ def load_and_predict(path_like_element): plot.savefig(os.path.join(os.path.dirname(path_like_element), f'latent_space_{idx}.png')) -def viz_latent(prediction, title=f'Latent Space '): +def viz_latent(prediction): if prediction.shape[-1] <= 1: raise ValueError('How did this happen?') elif prediction.shape[-1] == 2: