aboutsummaryrefslogtreecommitdiff
path: root/code/sunlab/sunflow/models/encoder_discriminator.py
blob: 5efb6afca0cc807e94e8e951a5db4578d0032905 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
class EncoderDiscriminator:
    """# EncoderDiscriminator Model

    Constructs an encoder-discriminator model"""

    def __init__(self, model_base_directory):
        """# EncoderDiscriminator Model Initialization

        - model_base_directory: The base folder directory where the model will
        be saved/ loaded"""
        self.model_base_directory = model_base_directory

    def init(self, encoder, discriminator):
        """# Initialize a EncoderDiscriminator

        - encoder: The encoder to use
        - discriminator: The discriminator to use"""
        from tensorflow import keras

        self.load_parameters()
        self.model = keras.models.Sequential()
        self.model.add(encoder.model)
        self.model.add(discriminator.model)
        self.model._name = "EncoderDiscriminator"
        return self

    def load(self):
        """# Load an existing EncoderDiscriminator"""
        from os import listdir

        if "encoder_discriminator.keras" not in listdir(
            f"{self.model_base_directory}/portable/"
        ):
            return None
        import tensorflow as tf

        self.model = tf.keras.models.load_model(
            f"{self.model_base_directory}/portable/encoder_discriminator" + ".keras",
            compile=False,
        )
        self.model._name = "EncoderDiscriminator"
        return self

    def save(self, overwrite=False):
        """# Save the current EncoderDiscriminator

        - Overwrite: overwrite any existing encoder_discriminator that has been
        saved"""
        from os import listdir

        if overwrite:
            self.model.save(
                f"{self.model_base_directory}/portable/encoder_discriminator" + ".keras"
            )
            return True
        if "encoder_discriminator.keras" in listdir(
            f"{self.model_base_directory}/portable/"
        ):
            return False
        self.model.save(
            f"{self.model_base_directory}/portable/encoder_discriminator" + ".keras"
        )
        return True

    def load_parameters(self):
        """# Load EncoderDiscriminator Model Parameters from File
        The file needs to have the following parameters defined:
         - data_size: int
         - autoencoder_layer_size: int
         - latent_size: int
         - autoencoder_depth: int
         - dropout: float (set to 0. if you don't want a dropout layer)
         - use_leaky_relu: boolean"""
        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.layer_size = parameters["autoencoder_layer_size"]
        self.latent_size = parameters["latent_size"]
        self.depth = parameters["autoencoder_depth"]
        self.dropout = parameters["dropout"]
        self.use_leaky_relu = parameters["use_leaky_relu"]

    def summary(self):
        """# Returns the summary of the EncoderDiscriminator model"""
        return self.model.summary()

    def __call__(self, *args, **kwargs):
        """# Callable

        When calling the encoder_discriminator class, return the model's
        output"""
        return self.model(*args, **kwargs)