hypermodels.py 5.87 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
            "units_conv", min_value=32, max_value=256, step=32, default=256,
33
        )
lucas_miranda's avatar
lucas_miranda committed
34
        lstm_units_1 = hp.Int(
35
            "units_lstm", min_value=128, max_value=512, step=32, default=256,
36
        )
lucas_miranda's avatar
lucas_miranda committed
37
        dense_2 = hp.Int(
38
            "units_dense2", min_value=32, max_value=256, step=32, default=64,
39
        )
lucas_miranda's avatar
lucas_miranda committed
40
        dropout_rate = hp.Float(
41
            "dropout_rate", min_value=0.0, max_value=0.5, default=0.25, step=0.05,
42
        )
43
        encoding = hp.Int("encoding", min_value=16, max_value=64, step=8, default=24,)
44

45
46
47
48
49
50
51
52
53
54
        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
55
56
57
58
59
60
61
62
63
64
        # 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]
65
66

        return model
67
68


69
class SEQ_2_SEQ_GMVAE(HyperModel):
70
71
    """Hyperparameter tuning pipeline for deepof.models.SEQ_2_SEQ_GMVAE"""

72
    def __init__(
lucas_miranda's avatar
lucas_miranda committed
73
        self,
74
75
76
        input_shape: tuple,
        entropy_reg_weight: float = 0.0,
        huber_delta: float = 1.0,
77
        kl_warmup_epochs: int = 0,
78
79
        learn_rate: float = 1e-3,
        loss: str = "ELBO+MMD",
80
        mmd_warmup_epochs: int = 0,
81
        number_of_components: int = 10,
82
        overlap_loss: float = False,
83
84
85
        phenotype_predictor: float = 0.0,
        predictor: float = 0.0,
        prior: str = "standard_normal",
86
    ):
87
        super().__init__()
88
        self.input_shape = input_shape
lucas_miranda's avatar
lucas_miranda committed
89
90
        self.entropy_reg_weight = entropy_reg_weight
        self.huber_delta = huber_delta
91
        self.kl_warmup_epochs = kl_warmup_epochs
92
        self.learn_rate = learn_rate
93
        self.loss = loss
94
        self.mmd_warmup_epochs = mmd_warmup_epochs
95
        self.number_of_components = number_of_components
lucas_miranda's avatar
lucas_miranda committed
96
        self.overlap_loss = overlap_loss
97
        self.pheno_class = phenotype_predictor
98
        self.predictor = predictor
lucas_miranda's avatar
lucas_miranda committed
99
        self.prior = prior
100

101
        assert (
lucas_miranda's avatar
lucas_miranda committed
102
            "ELBO" in self.loss or "MMD" in self.loss
103
        ), "loss must be one of ELBO, MMD or ELBO+MMD (default)"
104

105
    def get_hparams(self, hp):
106
        """Retrieve hyperparameters to tune"""
107

108
        # Architectural hyperparameters
lucas_miranda's avatar
lucas_miranda committed
109
        clipvalue = hp.Float(
110
            "clipvalue", min_value=0.0, max_value=1.0, default=0.5, sampling="Linear"
lucas_miranda's avatar
lucas_miranda committed
111
        )
lucas_miranda's avatar
lucas_miranda committed
112
        conv_filters = hp.Int(
lucas_miranda's avatar
lucas_miranda committed
113
            "units_conv", min_value=128, max_value=160, step=16, default=128,
114
        )
lucas_miranda's avatar
lucas_miranda committed
115
        dense_2 = hp.Int(
lucas_miranda's avatar
lucas_miranda committed
116
            "units_dense2", min_value=120, max_value=180, step=10, default=150,
117
        )
lucas_miranda's avatar
lucas_miranda committed
118
        dense_activation = hp.Choice("dense_activation", values=["elu", "relu"])
119
        dense_layers_per_branch = hp.Int("dense_layers_per_branch", min_value=1, max_value=3, default=1)
lucas_miranda's avatar
lucas_miranda committed
120
        dropout_rate = hp.Float(
121
122
            "dropout_rate",
            min_value=0.0,
lucas_miranda's avatar
lucas_miranda committed
123
124
            max_value=0.15,
            default=0.0,
125
126
            sampling="linear",
        )
lucas_miranda's avatar
lucas_miranda committed
127
128
129
130
        encoding = 16
        k = self.number_of_components
        lstm_units_1 = hp.Int(
            "units_lstm", min_value=300, max_value=350, step=10, default=320,
lucas_miranda's avatar
lucas_miranda committed
131
        )
132

133
        return (
lucas_miranda's avatar
lucas_miranda committed
134
            clipvalue,
135
136
            conv_filters,
            dense_2,
lucas_miranda's avatar
lucas_miranda committed
137
            dense_activation,
138
            dense_layers_per_branch,
139
140
141
            dropout_rate,
            encoding,
            k,
lucas_miranda's avatar
lucas_miranda committed
142
            lstm_units_1,
143
        )
144
145
146
147
148

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

        # Hyperparameters to tune
149
        (
lucas_miranda's avatar
lucas_miranda committed
150
            clipvalue,
151
152
            conv_filters,
            dense_2,
lucas_miranda's avatar
lucas_miranda committed
153
            dense_activation,
154
            dense_layers_per_branch,
155
156
            dropout_rate,
            encoding,
157
            k,
lucas_miranda's avatar
lucas_miranda committed
158
            lstm_units_1,
159
        ) = self.get_hparams(hp)
160

lucas_miranda's avatar
lucas_miranda committed
161
        gmvaep, kl_warmup_callback, mmd_warmup_callback = deepof.models.SEQ_2_SEQ_GMVAE(
lucas_miranda's avatar
lucas_miranda committed
162
            architecture_hparams={
lucas_miranda's avatar
lucas_miranda committed
163
164
                "clipvalue": clipvalue,
                "dense_activation": dense_activation,
lucas_miranda's avatar
lucas_miranda committed
165
                "dropout_rate": dropout_rate,
166
                "dense_layers_per_branch": dense_layers_per_branch,
lucas_miranda's avatar
lucas_miranda committed
167
                "encoding": encoding,
lucas_miranda's avatar
lucas_miranda committed
168
169
170
                "units_conv": conv_filters,
                "units_dense_2": dense_2,
                "units_lstm": lstm_units_1,
lucas_miranda's avatar
lucas_miranda committed
171
172
173
            },
            entropy_reg_weight=self.entropy_reg_weight,
            huber_delta=self.huber_delta,
174
            kl_warmup_epochs=self.kl_warmup_epochs,
lucas_miranda's avatar
lucas_miranda committed
175
            loss=self.loss,
176
            mmd_warmup_epochs=self.mmd_warmup_epochs,
177
            number_of_components=k,
lucas_miranda's avatar
lucas_miranda committed
178
            overlap_loss=self.overlap_loss,
179
            phenotype_prediction=self.pheno_class,
lucas_miranda's avatar
lucas_miranda committed
180
            predictor=self.predictor,
lucas_miranda's avatar
lucas_miranda committed
181
182
        ).build(self.input_shape)[3:]

lucas_miranda's avatar
lucas_miranda committed
183
        return gmvaep
184
185


lucas_miranda's avatar
lucas_miranda committed
186
187
188
# 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
189
#    - Number of dense layers close to the latent space as a hyperparameter (!)