hypermodels.py 5.65 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
14
15
16
17
import tensorflow_probability as tfp

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


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

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

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

lucas_miranda's avatar
lucas_miranda committed
31
        conv_filters = hp.Int(
32
33
34
35
36
            "units_conv",
            min_value=32,
            max_value=256,
            step=32,
            default=256,
37
        )
lucas_miranda's avatar
lucas_miranda committed
38
        lstm_units_1 = hp.Int(
39
40
41
42
43
            "units_lstm",
            min_value=128,
            max_value=512,
            step=32,
            default=256,
44
        )
lucas_miranda's avatar
lucas_miranda committed
45
        dense_2 = hp.Int(
46
47
48
49
50
            "units_dense2",
            min_value=32,
            max_value=256,
            step=32,
            default=64,
51
        )
lucas_miranda's avatar
lucas_miranda committed
52
        dropout_rate = hp.Float(
53
54
55
56
57
58
59
60
61
62
63
64
            "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,
65
66
        )

67
68
69
70
71
72
73
74
75
76
        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
77
78
79
80
81
82
83
84
85
86
        # 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]
87
88

        return model
89
90


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

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

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

125
    def get_hparams(self, hp):
126
        """Retrieve hyperparameters to tune"""
127

128
        # Architectural hyperparameters
lucas_miranda's avatar
lucas_miranda committed
129
130
        bidirectional_merge = "ave"
        clipvalue = 1.0
131
132
        conv_filters = hp.Int("conv_units", min_value=32, max_value=512, step=32)
        dense_2 = hp.Int("dense_units", min_value=32, max_value=512, step=32)
lucas_miranda's avatar
lucas_miranda committed
133
        dense_activation = "relu"
134
135
136
137
138
139
        dense_layers_per_branch = hp.Int(
            "dense_layers", min_value=1, max_value=3, step=1
        )
        dropout_rate = hp.Float(
            "dropout_rate", min_value=0.0, max_value=1.0, sampling="linear"
        )
lucas_miranda's avatar
lucas_miranda committed
140
        k = self.number_of_components
141
        lstm_units_1 = hp.Int("units", min_value=32, max_value=512, step=32)
142

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

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

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

171
        gmvaep = deepof.models.SEQ_2_SEQ_GMVAE(
lucas_miranda's avatar
lucas_miranda committed
172
            architecture_hparams={
173
                "bidirectional_merge": "ave",
lucas_miranda's avatar
lucas_miranda committed
174
175
                "clipvalue": clipvalue,
                "dense_activation": dense_activation,
176
                "dense_layers_per_branch": dense_layers_per_branch,
lucas_miranda's avatar
lucas_miranda committed
177
                "dropout_rate": dropout_rate,
lucas_miranda's avatar
lucas_miranda committed
178
179
180
                "units_conv": conv_filters,
                "units_dense_2": dense_2,
                "units_lstm": lstm_units_1,
lucas_miranda's avatar
lucas_miranda committed
181
            },
182
            encoding=self.encoding,
183
            kl_warmup_epochs=self.kl_warmup_epochs,
lucas_miranda's avatar
lucas_miranda committed
184
            loss=self.loss,
185
            mmd_warmup_epochs=self.mmd_warmup_epochs,
186
            number_of_components=k,
lucas_miranda's avatar
lucas_miranda committed
187
            overlap_loss=self.overlap_loss,
188
            phenotype_prediction=self.pheno_class,
lucas_miranda's avatar
lucas_miranda committed
189
            predictor=self.predictor,
190
        ).build(self.input_shape)[-1]
lucas_miranda's avatar
lucas_miranda committed
191

lucas_miranda's avatar
lucas_miranda committed
192
        return gmvaep
193
194


lucas_miranda's avatar
lucas_miranda committed
195
196
197
# 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
198
#    - Number of dense layers close to the latent space as a hyperparameter (!)