hypermodels.py 5.79 KB
Newer Older
lucas_miranda's avatar
lucas_miranda committed
1
# @author lucasmiranda42
2
3
4
5
6
7
8
9
# encoding: utf-8
# module deepof

"""

keras hypermodels for hyperparameter tuning of deep autoencoders

"""
lucas_miranda's avatar
lucas_miranda committed
10

11
from kerastuner import HyperModel
lucas_miranda's avatar
lucas_miranda committed
12
import deepof.models
13
import deepof.model_utils
lucas_miranda's avatar
lucas_miranda committed
14
import numpy as np
15
16
17
18
import tensorflow_probability as tfp

tfd = tfp.distributions
tfpl = tfp.layers
19
20


21
class SEQ_2_SEQ_AE(HyperModel):
22
23
    """Hyperparameter tuning pipeline for deepof.models.SEQ_2_SEQ_AE"""

24
    def __init__(self, input_shape):
25
        super().__init__()
26
27
        self.input_shape = input_shape

28
29
30
    @staticmethod
    def get_hparams(hp):
        """Retrieve hyperparameters to tune"""
31

lucas_miranda's avatar
lucas_miranda committed
32
        conv_filters = hp.Int(
33
34
35
36
37
            "units_conv",
            min_value=32,
            max_value=256,
            step=32,
            default=256,
38
        )
lucas_miranda's avatar
lucas_miranda committed
39
        lstm_units_1 = hp.Int(
40
41
42
43
44
            "units_lstm",
            min_value=128,
            max_value=512,
            step=32,
            default=256,
45
        )
lucas_miranda's avatar
lucas_miranda committed
46
        dense_2 = hp.Int(
47
48
49
50
51
            "units_dense2",
            min_value=32,
            max_value=256,
            step=32,
            default=64,
52
        )
lucas_miranda's avatar
lucas_miranda committed
53
        dropout_rate = hp.Float(
54
55
56
57
58
59
60
61
62
63
64
65
            "dropout_rate",
            min_value=0.0,
            max_value=0.5,
            default=0.25,
            step=0.05,
        )
        encoding = hp.Int(
            "encoding",
            min_value=16,
            max_value=64,
            step=8,
            default=24,
66
67
        )

68
69
70
71
72
73
74
75
76
77
        return conv_filters, lstm_units_1, dense_2, dropout_rate, encoding

    def build(self, hp):
        """Overrides Hypermodel's build method"""

        # HYPERPARAMETERS TO TUNE
        conv_filters, lstm_units_1, dense_2, dropout_rate, encoding = self.get_hparams(
            hp
        )

lucas_miranda's avatar
lucas_miranda committed
78
79
80
81
82
83
84
85
86
87
        # INSTANCIATED MODEL
        model = deepof.models.SEQ_2_SEQ_AE(
            architecture_hparams={
                "units_conv": conv_filters,
                "units_lstm": lstm_units_1,
                "units_dense_2": dense_2,
                "dropout_rate": dropout_rate,
                "encoding": encoding,
            }
        ).build(self.input_shape)[2]
88
89

        return model
90
91


92
class SEQ_2_SEQ_GMVAE(HyperModel):
93
94
    """Hyperparameter tuning pipeline for deepof.models.SEQ_2_SEQ_GMVAE"""

95
    def __init__(
lucas_miranda's avatar
lucas_miranda committed
96
        self,
97
98
99
        input_shape: tuple,
        entropy_reg_weight: float = 0.0,
        huber_delta: float = 1.0,
100
        kl_warmup_epochs: int = 0,
101
102
        learn_rate: float = 1e-3,
        loss: str = "ELBO+MMD",
103
        mmd_warmup_epochs: int = 0,
104
        number_of_components: int = 10,
105
        overlap_loss: float = False,
106
107
108
        phenotype_predictor: float = 0.0,
        predictor: float = 0.0,
        prior: str = "standard_normal",
109
    ):
110
        super().__init__()
111
        self.input_shape = input_shape
lucas_miranda's avatar
lucas_miranda committed
112
113
        self.entropy_reg_weight = entropy_reg_weight
        self.huber_delta = huber_delta
114
        self.kl_warmup_epochs = kl_warmup_epochs
115
        self.learn_rate = learn_rate
116
        self.loss = loss
117
        self.mmd_warmup_epochs = mmd_warmup_epochs
118
        self.number_of_components = number_of_components
lucas_miranda's avatar
lucas_miranda committed
119
        self.overlap_loss = overlap_loss
120
        self.pheno_class = phenotype_predictor
121
        self.predictor = predictor
lucas_miranda's avatar
lucas_miranda committed
122
        self.prior = prior
123

124
        assert (
lucas_miranda's avatar
lucas_miranda committed
125
            "ELBO" in self.loss or "MMD" in self.loss
126
        ), "loss must be one of ELBO, MMD or ELBO+MMD (default)"
127

128
    def get_hparams(self, hp):
129
        """Retrieve hyperparameters to tune"""
130

131
        # Architectural hyperparameters
lucas_miranda's avatar
lucas_miranda committed
132
133
        bidirectional_merge = hp.Choice(
            "bidirectional_merge", values=["sum", "mul", "concat", "ave"]
134
        )
lucas_miranda's avatar
lucas_miranda committed
135
        clipvalue = hp.Choice("clipvalue", values=[1.0, np.inf])
lucas_miranda's avatar
lucas_miranda committed
136
137
138
139
140
        conv_filters = 160
        dense_2 = 120
        dense_activation = "relu"
        dense_layers_per_branch = 1
        dropout_rate = 1e-3
lucas_miranda's avatar
lucas_miranda committed
141
142
        encoding = 16
        k = self.number_of_components
lucas_miranda's avatar
lucas_miranda committed
143
        lstm_units_1 = 300
144

145
        return (
lucas_miranda's avatar
lucas_miranda committed
146
            bidirectional_merge,
lucas_miranda's avatar
lucas_miranda committed
147
            clipvalue,
148
149
            conv_filters,
            dense_2,
lucas_miranda's avatar
lucas_miranda committed
150
            dense_activation,
151
            dense_layers_per_branch,
152
153
154
            dropout_rate,
            encoding,
            k,
lucas_miranda's avatar
lucas_miranda committed
155
            lstm_units_1,
156
        )
157
158
159
160
161

    def build(self, hp):
        """Overrides Hypermodel's build method"""

        # Hyperparameters to tune
162
        (
lucas_miranda's avatar
lucas_miranda committed
163
            bidirectional_merge,
lucas_miranda's avatar
lucas_miranda committed
164
            clipvalue,
165
166
            conv_filters,
            dense_2,
lucas_miranda's avatar
lucas_miranda committed
167
            dense_activation,
168
            dense_layers_per_branch,
169
170
            dropout_rate,
            encoding,
171
            k,
lucas_miranda's avatar
lucas_miranda committed
172
            lstm_units_1,
173
        ) = self.get_hparams(hp)
174

lucas_miranda's avatar
lucas_miranda committed
175
        gmvaep, kl_warmup_callback, mmd_warmup_callback = deepof.models.SEQ_2_SEQ_GMVAE(
lucas_miranda's avatar
lucas_miranda committed
176
            architecture_hparams={
lucas_miranda's avatar
lucas_miranda committed
177
                "bidirectional_merge": "concat",
lucas_miranda's avatar
lucas_miranda committed
178
179
                "clipvalue": clipvalue,
                "dense_activation": dense_activation,
180
                "dense_layers_per_branch": dense_layers_per_branch,
lucas_miranda's avatar
lucas_miranda committed
181
                "dropout_rate": dropout_rate,
lucas_miranda's avatar
lucas_miranda committed
182
                "encoding": encoding,
lucas_miranda's avatar
lucas_miranda committed
183
184
185
                "units_conv": conv_filters,
                "units_dense_2": dense_2,
                "units_lstm": lstm_units_1,
lucas_miranda's avatar
lucas_miranda committed
186
187
188
            },
            entropy_reg_weight=self.entropy_reg_weight,
            huber_delta=self.huber_delta,
189
            kl_warmup_epochs=self.kl_warmup_epochs,
lucas_miranda's avatar
lucas_miranda committed
190
            loss=self.loss,
191
            mmd_warmup_epochs=self.mmd_warmup_epochs,
192
            number_of_components=k,
lucas_miranda's avatar
lucas_miranda committed
193
            overlap_loss=self.overlap_loss,
194
            phenotype_prediction=self.pheno_class,
lucas_miranda's avatar
lucas_miranda committed
195
            predictor=self.predictor,
lucas_miranda's avatar
lucas_miranda committed
196
197
        ).build(self.input_shape)[3:]

lucas_miranda's avatar
lucas_miranda committed
198
        return gmvaep
199
200


lucas_miranda's avatar
lucas_miranda committed
201
202
203
# TODO:
#    - We can add as many parameters as we want to the hypermodel!
#    with this implementation, predictor, warmup, loss and even number of components can be tuned using BayOpt
lucas_miranda's avatar
lucas_miranda committed
204
#    - Number of dense layers close to the latent space as a hyperparameter (!)