steffen/ae_toolbox_torch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Visualization approach n

Browse Source
This commit is contained in:
Si11ium
2019-09-29 09:37:30 +02:00
parent 1386cdfd33
commit a70c9b7fef
15 changed files with 652 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
run_models.py
View File

@ -1,30 +1,32 @@
from torch.distributions import Normal
from networks.auto_encoder import *
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 pytorch_lightning import Trainer
from test_tube import Experiment
import os
from argparse import Namespace
from argparse import ArgumentParser
from distutils.util import strtobool
from networks.auto_encoder import AutoEncoder, AutoEncoder_LO
from networks.variational_auto_encoder import VariationalAE, VAE_LO
from networks.adverserial_auto_encoder import AdversarialAE_LO, AdversarialAE
from networks.seperating_adversarial_auto_encoder import SeperatingAAE, SeparatingAAE_LO, SuperSeperatingAAE
from networks.modules import LightningModule
from pytorch_lightning import Trainer
from test_tube import Experiment
args = ArgumentParser()
args.add_argument('--step', default=6)
args.add_argument('--step', default=5)
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('--latent_dim', default=2)
args.add_argument('--model', default='VAE_Model')
args.add_argument('--refresh', type=strtobool, default=False)
# ToDo: How to implement this better?
# other_classes = [AutoEncoder, AutoEncoderLightningOverrides]
class Model(AutoEncoderLightningOverrides, LightningModule):
class AE_Model(AutoEncoder_LO, LightningModule):
def __init__(self, parameters):
assert all([x in parameters for x in ['step', 'size', 'latent_dim', 'features']])
@ -32,11 +34,23 @@ class Model(AutoEncoderLightningOverrides, LightningModule):
self.latent_dim = parameters.latent_dim
self.features = parameters.features
self.step = parameters.step
super(Model, self).__init__()
super(AE_Model, self).__init__()
self.network = AutoEncoder(self.latent_dim, self.features)
class AdversarialModel(AdversarialAELightningOverrides, LightningModule):
class VAE_Model(VAE_LO, LightningModule):
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
self.features = parameters.features
self.step = parameters.step
super(VAE_Model, self).__init__()
self.network = VariationalAE(self.latent_dim, self.features)
class AAE_Model(AdversarialAE_LO, LightningModule):
def __init__(self, parameters: Namespace):
assert all([x in parameters for x in ['step', 'size', 'latent_dim', 'features']])
@ -44,13 +58,13 @@ class AdversarialModel(AdversarialAELightningOverrides, LightningModule):
self.latent_dim = parameters.latent_dim
self.features = parameters.features
self.step = parameters.step
super(AdversarialModel, self).__init__()
super(AAE_Model, self).__init__()
self.normal = Normal(0, 1)
self.network = AdversarialAutoEncoder(self.latent_dim, self.features)
self.network = AdversarialAE(self.latent_dim, self.features)
pass
class SeparatingAdversarialModel(SeparatingAdversarialAELightningOverrides, LightningModule):
class SAAE_Model(SeparatingAAE_LO, LightningModule):
def __init__(self, parameters: Namespace):
assert all([x in parameters for x in ['step', 'size', 'latent_dim', 'features']])
@ -58,9 +72,23 @@ class SeparatingAdversarialModel(SeparatingAdversarialAELightningOverrides, Ligh
self.latent_dim = parameters.latent_dim
self.features = parameters.features
self.step = parameters.step
super(SeparatingAdversarialModel, self).__init__()
super(SAAE_Model, self).__init__()
self.normal = Normal(0, 1)
self.network = SeperatingAdversarialAutoEncoder(self.latent_dim, self.features)
self.network = SeperatingAAE(self.latent_dim, self.features)
pass
class SSAAE_Model(SeparatingAAE_LO, LightningModule):
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
self.features = parameters.features
self.step = parameters.step
super(SSAAE_Model, self).__init__()
self.normal = Normal(0, 1)
self.network = SuperSeperatingAAE(self.latent_dim, self.features)
pass
@ -84,8 +112,13 @@ if __name__ == '__main__':
period=4
)
trainer = Trainer(experiment=exp, max_nb_epochs=250, gpus=[0],
add_log_row_interval=1000, checkpoint_callback=checkpoint_callback)
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'))