From aa802cb2be008397c4909351566b6805d542d5c6 Mon Sep 17 00:00:00 2001 From: Si11ium Date: Sun, 29 Sep 2019 11:52:39 +0200 Subject: [PATCH] Future Prediction Training as Parameter --- run_models.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/run_models.py b/run_models.py index 88ad3b1..0a78648 100644 --- a/run_models.py +++ b/run_models.py @@ -24,6 +24,7 @@ args.add_argument('--size', default=9) args.add_argument('--latent_dim', default=2) args.add_argument('--model', default='AE_Model') args.add_argument('--refresh', type=strtobool, default=False) +args.add_argument('--future_predictions', type=strtobool, default=False) class AE_Model(AutoEncoder_LO, LightningModule): @@ -34,7 +35,7 @@ class AE_Model(AutoEncoder_LO, LightningModule): self.latent_dim = parameters.latent_dim self.features = parameters.features self.step = parameters.step - super(AE_Model, self).__init__() + super(AE_Model, self).__init__(train_on_predictions=parameters.future_predictions) self.network = AutoEncoder(self.latent_dim, self.features) @@ -46,7 +47,7 @@ class VAE_Model(VAE_LO, LightningModule): self.latent_dim = parameters.latent_dim self.features = parameters.features self.step = parameters.step - super(VAE_Model, self).__init__() + super(VAE_Model, self).__init__(train_on_predictions=parameters.future_predictions) self.network = VariationalAE(self.latent_dim, self.features) @@ -58,7 +59,7 @@ class AAE_Model(AdversarialAE_LO, LightningModule): self.latent_dim = parameters.latent_dim self.features = parameters.features self.step = parameters.step - super(AAE_Model, self).__init__() + super(AAE_Model, self).__init__(train_on_predictions=parameters.future_predictions) self.normal = Normal(0, 1) self.network = AdversarialAE(self.latent_dim, self.features) pass @@ -72,7 +73,7 @@ class SAAE_Model(SeparatingAAE_LO, LightningModule): self.latent_dim = parameters.latent_dim self.features = parameters.features self.step = parameters.step - super(SAAE_Model, self).__init__() + super(SAAE_Model, self).__init__(train_on_predictions=parameters.future_predictions) self.normal = Normal(0, 1) self.network = SeperatingAAE(self.latent_dim, self.features) pass