model_utils.py 5.47 KB
Newer Older
1
2
3
# @author lucasmiranda42

from keras import backend as K
4
from sklearn.metrics import silhouette_score
5
6
7
from tensorflow.keras.constraints import Constraint
from tensorflow.keras.layers import Layer
import tensorflow as tf
8
import tensorflow_probability as tfp
9

10
tfd = tfp.distributions
11
tfpl = tfp.layers
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

# Helper functions
def compute_kernel(x, y):
    x_size = K.shape(x)[0]
    y_size = K.shape(y)[0]
    dim = K.shape(x)[1]
    tiled_x = K.tile(K.reshape(x, K.stack([x_size, 1, dim])), K.stack([1, y_size, 1]))
    tiled_y = K.tile(K.reshape(y, K.stack([1, y_size, dim])), K.stack([x_size, 1, 1]))
    return K.exp(
        -tf.reduce_mean(K.square(tiled_x - tiled_y), axis=2) / K.cast(dim, tf.float32)
    )


def compute_mmd(x, y):
    x_kernel = compute_kernel(x, x)
    y_kernel = compute_kernel(y, y)
    xy_kernel = compute_kernel(x, y)
    return (
        tf.reduce_mean(x_kernel)
        + tf.reduce_mean(y_kernel)
        - 2 * tf.reduce_mean(xy_kernel)
    )


# Custom layers for efficiency/losses
class DenseTranspose(Layer):
    def __init__(self, dense, output_dim, activation=None, **kwargs):
        self.dense = dense
        self.output_dim = output_dim
        self.activation = tf.keras.activations.get(activation)
        super().__init__(**kwargs)

    def get_config(self):
        config = super().get_config().copy()
        config.update(
            {
                "dense": self.dense,
                "output_dim": self.output_dim,
                "activation": self.activation,
            }
        )
        return config

    def build(self, batch_input_shape):
        self.biases = self.add_weight(
            name="bias", shape=[self.dense.input_shape[-1]], initializer="zeros"
        )
        super().build(batch_input_shape)

    def call(self, inputs, **kwargs):
        z = tf.matmul(inputs, self.dense.weights[0], transpose_b=True)
        return self.activation(z + self.biases)

    def compute_output_shape(self, input_shape):
        return input_shape[0], self.output_dim


class UncorrelatedFeaturesConstraint(Constraint):
    def __init__(self, encoding_dim, weightage=1.0):
        self.encoding_dim = encoding_dim
        self.weightage = weightage

    def get_config(self):

        config = super().get_config().copy()
        config.update(
            {"encoding_dim": self.encoding_dim, "weightage": self.weightage,}
        )
        return config

    def get_covariance(self, x):
        x_centered_list = []

        for i in range(self.encoding_dim):
            x_centered_list.append(x[:, i] - K.mean(x[:, i]))

        x_centered = tf.stack(x_centered_list)
        covariance = K.dot(x_centered, K.transpose(x_centered)) / tf.cast(
            x_centered.get_shape()[0], tf.float32
        )

        return covariance

    # Constraint penalty
    def uncorrelated_feature(self, x):
        if self.encoding_dim <= 1:
            return 0.0
        else:
            output = K.sum(
                K.square(
                    self.covariance
                    - tf.math.multiply(self.covariance, K.eye(self.encoding_dim))
                )
            )
            return output

    def __call__(self, x):
        self.covariance = self.get_covariance(x)
        return self.weightage * self.uncorrelated_feature(x)


113
114
class KLDivergenceLayer(tfpl.KLDivergenceAddLoss):
    def __init__(self, *args, **kwargs):
115
116
117
        self.is_placeholder = True
        super(KLDivergenceLayer, self).__init__(*args, **kwargs)

118
119
120
121
122
    def call(self, distribution_a):
        kl_batch = self._regularizer(distribution_a)
        self.add_loss(kl_batch, inputs=[distribution_a])
        self.add_metric(
            kl_batch, aggregation="mean", name="kl_divergence",
123
        )
124
        self.add_metric(self._regularizer._weight, aggregation="mean", name="kl_rate")
125

126
        return distribution_a
127
128
129
130
131
132
133


class MMDiscrepancyLayer(Layer):
    """ Identity transform layer that adds MM discrepancy
    to the final model loss.
    """

134
    def __init__(self, prior, beta=1.0, *args, **kwargs):
135
        self.is_placeholder = True
136
        self.beta = beta
137
        self.prior = prior
138
139
        super(MMDiscrepancyLayer, self).__init__(*args, **kwargs)

140
141
142
    def get_config(self):
        config = super().get_config().copy()
        config.update({"beta": self.beta})
143
        config.update({"prior": self.prior})
144
145
        return config

146
    def call(self, z, **kwargs):
147
        true_samples = self.prior.sample(1)
148
149
        mmd_batch = self.beta * compute_mmd(true_samples, z)
        self.add_loss(K.mean(mmd_batch), inputs=z)
150
        self.add_metric(mmd_batch, aggregation="mean", name="mmd")
151
        self.add_metric(self.beta, aggregation="mean", name="mmd_rate")
152
153

        return z
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173


class Latent_space_control(Layer):
    """ Identity layer that adds latent space and clustering stats
     to the metrics compiled by the model
     """

    def __init__(self, *args, **kwargs):
        super(Latent_space_control, self).__init__(*args, **kwargs)

    def call(self, z, z_gauss, z_cat, **kwargs):

        # Adds metric that monitors dead neurons in the latent space
        self.add_metric(
            tf.math.zero_fraction(z_gauss), aggregation="mean", name="dead_neurons"
        )

        # Adds Silhouette score controling overlap between clusters
        hard_labels = tf.math.argmax(z_cat, axis=1)
        silhouette = tf.numpy_function(silhouette_score, [z, hard_labels], tf.float32)
174
        self.add_loss(-K.mean(silhouette), inputs=[z, hard_labels])
175
176
177
        self.add_metric(silhouette, aggregation="mean", name="silhouette")

        return z