Initial commit

1 files changed, 344 insertions, 0 deletions

diff --git a/code/sunlab/sunflow/models/adversarial_autoencoder.py b/code/sunlab/sunflow/models/adversarial_autoencoder.py
new file mode 100644
index 0000000..4cbb2f8
--- /dev/null
+++ b/code/sunlab/sunflow/models/adversarial_autoencoder.py
@@ -0,0 +1,344 @@
+from sunlab.common.data.dataset import Dataset
+from sunlab.common.scaler.adversarial_scaler import AdversarialScaler
+from sunlab.common.distribution.adversarial_distribution import AdversarialDistribution
+from .encoder import Encoder
+from .decoder import Decoder
+from .discriminator import Discriminator
+from .encoder_discriminator import EncoderDiscriminator
+from .autoencoder import Autoencoder
+from tensorflow.keras import optimizers, metrics, losses
+import tensorflow as tf
+from numpy import ones, zeros, float32, NaN
+
+
+class AdversarialAutoencoder:
+    """# Adversarial Autoencoder
+    - distribution: The distribution used by the adversary to learn on"""
+
+    def __init__(
+        self,
+        model_base_directory,
+        distribution: AdversarialDistribution or None = None,
+        scaler: AdversarialScaler or None = None,
+    ):
+        """# Adversarial Autoencoder Model Initialization
+
+        - model_base_directory: The base folder directory where the model will
+        be saved/ loaded
+        - distribution: The distribution the adversary will use
+        - scaler: The scaling function the model will assume on the data"""
+        self.model_base_directory = model_base_directory
+        if distribution is not None:
+            self.distribution = distribution
+        else:
+            self.distribution = None
+        if scaler is not None:
+            self.scaler = scaler(self.model_base_directory)
+        else:
+            self.scaler = None
+
+    def init(
+        self,
+        data=None,
+        data_size=13,
+        autoencoder_layer_size=16,
+        adversary_layer_size=8,
+        latent_size=2,
+        autoencoder_depth=2,
+        dropout=0.0,
+        use_leaky_relu=False,
+        **kwargs,
+    ):
+        """# Initialize AAE model parameters
+        - data_size: int
+        - autoencoder_layer_size: int
+        - adversary_layer_size: int
+        - latent_size: int
+        - autoencoder_depth: int
+        - dropout: float
+        - use_leaky_relu: boolean"""
+        self.data_size = data_size
+        self.autoencoder_layer_size = autoencoder_layer_size
+        self.adversary_layer_size = adversary_layer_size
+        self.latent_size = latent_size
+        self.autoencoder_depth = autoencoder_depth
+        self.dropout = dropout
+        self.use_leaky_relu = use_leaky_relu
+        self.save_parameters()
+        self.encoder = Encoder(self.model_base_directory).init()
+        self.decoder = Decoder(self.model_base_directory).init()
+        self.autoencoder = Autoencoder(self.model_base_directory).init(
+            self.encoder, self.decoder
+        )
+        self.discriminator = Discriminator(self.model_base_directory).init()
+        self.encoder_discriminator = EncoderDiscriminator(
+            self.model_base_directory
+        ).init(self.encoder, self.discriminator)
+        if self.distribution is not None:
+            self.distribution = self.distribution(self.latent_size)
+        if (data is not None) and (self.scaler is not None):
+            self.scaler = self.scaler.init(data)
+        self.init_optimizers_and_metrics(**kwargs)
+        return self
+
+    def init_optimizers_and_metrics(
+        self,
+        optimizer=optimizers.Adam,
+        ae_metric=metrics.MeanAbsoluteError,
+        adv_metric=metrics.BinaryCrossentropy,
+        ae_lr=7e-4,
+        adv_lr=3e-4,
+        loss_fn=losses.BinaryCrossentropy,
+        **kwargs,
+    ):
+        """# Set the optimizer, loss function, and metrics"""
+        self.ae_optimizer = optimizer(learning_rate=ae_lr)
+        self.adv_optimizer = optimizer(learning_rate=adv_lr)
+        self.gan_optimizer = optimizer(learning_rate=adv_lr)
+        self.train_ae_metric = ae_metric()
+        self.val_ae_metric = ae_metric()
+        self.train_adv_metric = adv_metric()
+        self.val_adv_metric = adv_metric()
+        self.train_gan_metric = adv_metric()
+        self.val_gan_metric = adv_metric()
+        self.loss_fn = loss_fn()
+
+    def load(self):
+        """# Load the models from their respective files"""
+        self.load_parameters()
+        self.encoder = Encoder(self.model_base_directory).load()
+        self.decoder = Decoder(self.model_base_directory).load()
+        self.autoencoder = Autoencoder(self.model_base_directory).load()
+        self.discriminator = Discriminator(self.model_base_directory).load()
+        self.encoder_discriminator = EncoderDiscriminator(
+            self.model_base_directory
+        ).load()
+        if self.scaler is not None:
+            self.scaler = self.scaler.load()
+        return self
+
+    def save(self, overwrite=False):
+        """# Save each model in the AAE"""
+        self.encoder.save(overwrite=overwrite)
+        self.decoder.save(overwrite=overwrite)
+        self.autoencoder.save(overwrite=overwrite)
+        self.discriminator.save(overwrite=overwrite)
+        self.encoder_discriminator.save(overwrite=overwrite)
+        if self.scaler is not None:
+            self.scaler.save()
+
+    def save_parameters(self):
+        """# Save the AAE parameters in a file"""
+        from pickle import dump
+        from os import makedirs
+
+        makedirs(self.model_base_directory + "/portable/", exist_ok=True)
+        parameters = {
+            "data_size": self.data_size,
+            "autoencoder_layer_size": self.autoencoder_layer_size,
+            "adversary_layer_size": self.adversary_layer_size,
+            "latent_size": self.latent_size,
+            "autoencoder_depth": self.autoencoder_depth,
+            "dropout": self.dropout,
+            "use_leaky_relu": self.use_leaky_relu,
+        }
+        with open(
+            f"{self.model_base_directory}/portable/model_parameters.pkl", "wb"
+        ) as phandle:
+            dump(parameters, phandle)
+
+    def load_parameters(self):
+        """# Load the AAE parameters from a file"""
+        from pickle import load
+
+        with open(
+            f"{self.model_base_directory}/portable/model_parameters.pkl", "rb"
+        ) as phandle:
+            parameters = load(phandle)
+        self.data_size = parameters["data_size"]
+        self.autoencoder_layer_size = parameters["autoencoder_layer_size"]
+        self.adversary_layer_size = parameters["adversary_layer_size"]
+        self.latent_size = parameters["latent_size"]
+        self.autoencoder_depth = parameters["autoencoder_depth"]
+        self.dropout = parameters["dropout"]
+        self.use_leaky_relu = parameters["use_leaky_relu"]
+        return parameters
+
+    def summary(self):
+        """# Summarize each model in the AAE"""
+        self.encoder.summary()
+        self.decoder.summary()
+        self.autoencoder.summary()
+        self.discriminator.summary()
+        self.encoder_discriminator.summary()
+
+    @tf.function
+    def train_step(self, x, y):
+        """# Training Step
+
+        1. Train the Autoencoder
+        2. (If distribution is given) Train the discriminator
+        3. (If the distribution is given) Train the encoder_discriminator"""
+        # Autoencoder Training
+        with tf.GradientTape() as tape:
+            decoded_vector = self.autoencoder(x, training=True)
+            ae_loss_value = self.loss_fn(y, decoded_vector)
+        grads = tape.gradient(ae_loss_value, self.autoencoder.model.trainable_weights)
+        self.ae_optimizer.apply_gradients(
+            zip(grads, self.autoencoder.model.trainable_weights)
+        )
+        self.train_ae_metric.update_state(y, decoded_vector)
+        if self.distribution is not None:
+            # Adversary Trainig
+            with tf.GradientTape() as tape:
+                latent_vector = self.encoder(x)
+                fakepred = self.distribution(x.shape[0])
+                discbatch_x = tf.concat([latent_vector, fakepred], axis=0)
+                discbatch_y = tf.concat([zeros(x.shape[0]), ones(x.shape[0])], axis=0)
+                adversary_vector = self.discriminator(discbatch_x, training=True)
+                adv_loss_value = self.loss_fn(discbatch_y, adversary_vector)
+            grads = tape.gradient(
+                adv_loss_value, self.discriminator.model.trainable_weights
+            )
+            self.adv_optimizer.apply_gradients(
+                zip(grads, self.discriminator.model.trainable_weights)
+            )
+            self.train_adv_metric.update_state(discbatch_y, adversary_vector)
+            # Gan Training
+            with tf.GradientTape() as tape:
+                gan_vector = self.encoder_discriminator(x, training=True)
+                adv_vector = tf.convert_to_tensor(ones((x.shape[0], 1), dtype=float32))
+                gan_loss_value = self.loss_fn(gan_vector, adv_vector)
+            grads = tape.gradient(gan_loss_value, self.encoder.model.trainable_weights)
+            self.gan_optimizer.apply_gradients(
+                zip(grads, self.encoder.model.trainable_weights)
+            )
+            self.train_gan_metric.update_state(adv_vector, gan_vector)
+            return (ae_loss_value, adv_loss_value, gan_loss_value)
+        return (ae_loss_value, None, None)
+
+    @tf.function
+    def test_step(self, x, y):
+        """# Test Step - On validation data
+
+        1. Evaluate the Autoencoder
+        2. (If distribution is given) Evaluate the discriminator
+        3. (If the distribution is given) Evaluate the encoder_discriminator"""
+        val_decoded_vector = self.autoencoder(x, training=False)
+        self.val_ae_metric.update_state(y, val_decoded_vector)
+
+        if self.distribution is not None:
+            latent_vector = self.encoder(x)
+            fakepred = self.distribution(x.shape[0])
+            discbatch_x = tf.concat([latent_vector, fakepred], axis=0)
+            discbatch_y = tf.concat([zeros(x.shape[0]), ones(x.shape[0])], axis=0)
+            adversary_vector = self.discriminator(discbatch_x, training=False)
+            self.val_adv_metric.update_state(discbatch_y, adversary_vector)
+
+            gan_vector = self.encoder_discriminator(x, training=False)
+            self.val_gan_metric.update_state(ones(x.shape[0]), gan_vector)
+
+    # Garbage Collect at the end of each epoch
+    def on_epoch_end(self, _epoch, logs=None):
+        """# Cleanup environment to prevent memory leaks each epoch"""
+        import gc
+        from tensorflow.keras import backend as k
+
+        gc.collect()
+        k.clear_session()
+
+    def train(
+        self,
+        dataset: Dataset,
+        epoch_count: int = 1,
+        output=False,
+        output_freq=1,
+        fmt="%i[%.3f]: %.2e %.2e %.2e  %.2e %.2e %.2e",
+    ):
+        """# Train the model on a dataset
+
+         - dataset: ataset = Dataset to train the model on, which as the
+        training and validation iterators set up
+         - epoch_count: int = The number of epochs to train
+         - output: boolean =  Whether or not to output training information
+         - output_freq: int = The number of epochs between each output"""
+        from time import time
+        from numpy import array as narray
+
+        def fmtter(x):
+            return x if x is not None else -1
+
+        epoch_data = []
+        dataset.reset_iterators()
+
+        self.test_step(dataset.dataset, dataset.dataset)
+        val_ae = self.val_ae_metric.result()
+        val_adv = self.val_adv_metric.result()
+        val_gan = self.val_gan_metric.result()
+        self.val_ae_metric.reset_states()
+        self.val_adv_metric.reset_states()
+        self.val_gan_metric.reset_states()
+        print(
+            fmt
+            % (
+                0,
+                NaN,
+                val_ae,
+                fmtter(val_adv),
+                fmtter(val_gan),
+                NaN,
+                NaN,
+                NaN,
+            )
+        )
+        for epoch in range(epoch_count):
+            start_time = time()
+
+            for step, (x_batch_train, y_batch_train) in enumerate(dataset.training):
+                ae_lv, adv_lv, gan_lv = self.train_step(x_batch_train, x_batch_train)
+
+            train_ae = self.train_ae_metric.result()
+            train_adv = self.train_adv_metric.result()
+            train_gan = self.train_gan_metric.result()
+            self.train_ae_metric.reset_states()
+            self.train_adv_metric.reset_states()
+            self.train_gan_metric.reset_states()
+
+            for step, (x_batch_val, y_batch_val) in enumerate(dataset.validation):
+                self.test_step(x_batch_val, x_batch_val)
+
+            val_ae = self.val_ae_metric.result()
+            val_adv = self.val_adv_metric.result()
+            val_gan = self.val_gan_metric.result()
+            self.val_ae_metric.reset_states()
+            self.val_adv_metric.reset_states()
+            self.val_gan_metric.reset_states()
+
+            epoch_data.append(
+                (
+                    epoch,
+                    train_ae,
+                    val_ae,
+                    fmtter(train_adv),
+                    fmtter(val_adv),
+                    fmtter(train_gan),
+                    fmtter(val_gan),
+                )
+            )
+            if output and (epoch + 1) % output_freq == 0:
+                print(
+                    fmt
+                    % (
+                        epoch + 1,
+                        time() - start_time,
+                        train_ae,
+                        fmtter(train_adv),
+                        fmtter(train_gan),
+                        val_ae,
+                        fmtter(val_adv),
+                        fmtter(val_gan),
+                    )
+                )
+            self.on_epoch_end(epoch)
+            dataset.reset_iterators()
+        return narray(epoch_data)