Implemented deep clustering using a SEQ_2_SEQ_GMVAEP model

source/models.py

@@ -2,12 +2,13 @@
 
 from tensorflow.keras import backend as K
 from tensorflow.keras import Input, Model, Sequential
+from tensorflow.keras.activations import softplus
 from tensorflow.keras.callbacks import LambdaCallback
 from tensorflow.keras.constraints import UnitNorm
 from tensorflow.keras.initializers import he_uniform, Orthogonal
-from tensorflow.keras.layers import BatchNormalization, Bidirectional, Dense
-from tensorflow.keras.layers import Dropout, LSTM
-from tensorflow.keras.layers import RepeatVector, TimeDistributed
+from tensorflow.keras.layers import BatchNormalization, Bidirectional, Concatenate
+from tensorflow.keras.layers import Dense, Dropout, LSTM
+from tensorflow.keras.layers import RepeatVector, Reshape, TimeDistributed
 from tensorflow.keras.losses import Huber
 from tensorflow.keras.optimizers import Adam
 from source.model_utils import *
@@ -677,24 +678,15 @@ class SEQ_2_SEQ_MMVAEP:
             kernel_constraint=UnitNorm(axis=0),
             kernel_initializer=he_uniform(),
         )
-        Model_E5 = Dense(
-            self.ENCODING,
-            activation="relu",
-            kernel_constraint=UnitNorm(axis=1),
-            activity_regularizer=UncorrelatedFeaturesConstraint(3, weightage=1.0),
-            kernel_initializer=Orthogonal(),
-        )
 
         # Decoder layers
         Model_B1 = BatchNormalization()
         Model_B2 = BatchNormalization()
         Model_B3 = BatchNormalization()
         Model_B4 = BatchNormalization()
-        Model_B5 = BatchNormalization()
-        Model_D0 = DenseTranspose(
-            Model_E5, activation="relu", output_dim=self.ENCODING,
+        Model_D1 = Dense(
+            self.DENSE_2, activation="relu", kernel_initializer=he_uniform()
         )
-        Model_D1 = DenseTranspose(Model_E4, activation="relu", output_dim=self.DENSE_2,)
         Model_D2 = DenseTranspose(Model_E3, activation="relu", output_dim=self.DENSE_1,)
         Model_D3 = RepeatVector(self.input_shape[1])
         Model_D4 = Bidirectional(
@@ -727,15 +719,12 @@ class SEQ_2_SEQ_MMVAEP:
         encoder = Dropout(self.DROPOUT_RATE)(encoder)
         encoder = Model_E4(encoder)
         encoder = BatchNormalization()(encoder)
-        encoder = Model_E5(encoder)
-
-        # Map encoder to a categorical distribution over the components
-        zcat = Dense(self.number_of_components, activation="softmax")(encoder)
 
-        # Map encoder to a dense layer representing the parameters of
-        # the gaussian mixture latent space
-        zgauss = Dense(
-            tfpl.MixtureNormal.params_size(self.number_of_components, self.ENCODING),
+        z_cat = Dense(self.number_of_components, activation="softmax")(encoder)
+        z_gauss = Dense(
+            tfpl.IndependentNormal.params_size(
+                self.ENCODING * self.number_of_components
+            ),
             activation=None,
         )(encoder)
 
@@ -752,7 +741,23 @@ class SEQ_2_SEQ_MMVAEP:
                     )
                 )
 
-        z = tfpl.MixtureNormal(self.number_of_components, self.ENCODING)(zgauss)
+        z_gauss = Reshape([2 * self.ENCODING, self.number_of_components])(z_gauss)
+        z = tfpl.DistributionLambda(
+            lambda gauss: tfd.mixture.Mixture(
+                cat=tfd.categorical.Categorical(probs=gauss[0],),
+                components=[
+                    tfd.Independent(
+                        tfd.Normal(
+                            loc=gauss[1][..., : self.ENCODING, k],
+                            scale=softplus(gauss[1][..., self.ENCODING :, k]),
+                        ),
+                        reinterpreted_batch_ndims=1,
+                    )
+                    for k in range(self.number_of_components)
+                ],
+            ),
+            activity_regularizer=UncorrelatedFeaturesConstraint(3, weightage=1.0),
+        )([z_cat, z_gauss])
 
         if "ELBO" in self.loss:
             z = KLDivergenceLayer(self.prior, weight=kl_beta)(z)
@@ -772,29 +777,23 @@ class SEQ_2_SEQ_MMVAEP:
             z = MMDiscrepancyLayer(prior=self.prior, beta=mmd_beta)(z)
 
         # Define and instantiate generator
-        generator = Model_D0(z)
+        generator = Model_D1(z)
         generator = Model_B1(generator)
-        generator = Model_D1(generator)
-        generator = Model_B2(generator)
         generator = Model_D2(generator)
-        generator = Model_B3(generator)
+        generator = Model_B2(generator)
         generator = Model_D3(generator)
         generator = Model_D4(generator)
-        generator = Model_B4(generator)
+        generator = Model_B3(generator)
         generator = Model_D5(generator)
-        generator = Model_B5(generator)
+        generator = Model_B4(generator)
         x_decoded_mean = TimeDistributed(
             Dense(self.input_shape[2]), name="vaep_reconstruction"
         )(generator)
 
         # Define and instantiate predictor
-        predictor = Dense(
-            self.ENCODING, activation="relu", kernel_initializer=he_uniform()
-        )(z)
-        predictor = BatchNormalization()(predictor)
         predictor = Dense(
             self.DENSE_2, activation="relu", kernel_initializer=he_uniform()
-        )(predictor)
+        )(z)
         predictor = BatchNormalization()(predictor)
         predictor = Dense(
             self.DENSE_1, activation="relu", kernel_initializer=he_uniform()
@@ -825,23 +824,22 @@ class SEQ_2_SEQ_MMVAEP:
 
         # end-to-end autoencoder
         encoder = Model(x, z, name="SEQ_2_SEQ_VEncoder")
+        grouper = Model(x, z_cat, name="Deep_Gaussian_Mixture_clustering")
         gmvaep = Model(
-            inputs=x, outputs=[x_decoded_mean, x_predicted_mean], name="SEQ_2_SEQ_VAE"
+            inputs=x, outputs=[x_decoded_mean, x_predicted_mean], name="SEQ_2_SEQ_VAEP"
         )
 
         # Build generator as a separate entity
         g = Input(shape=self.ENCODING)
-        _generator = Model_D0(g)
+        _generator = Model_D1(g)
         _generator = Model_B1(_generator)
-        _generator = Model_D1(_generator)
-        _generator = Model_B2(_generator)
         _generator = Model_D2(_generator)
-        _generator = Model_B3(_generator)
+        _generator = Model_B2(_generator)
         _generator = Model_D3(_generator)
         _generator = Model_D4(_generator)
-        _generator = Model_B4(_generator)
+        _generator = Model_B3(_generator)
         _generator = Model_D5(_generator)
-        _generator = Model_B5(_generator)
+        _generator = Model_B4(_generator)
         _x_decoded_mean = TimeDistributed(Dense(self.input_shape[2]))(_generator)
         generator = Model(g, _x_decoded_mean, name="SEQ_2_SEQ_VGenerator")
 
@@ -853,7 +851,14 @@ class SEQ_2_SEQ_MMVAEP:
             loss=huber_loss, optimizer=Adam(lr=self.learn_rate,), metrics=["mae"],
         )
 
-        return encoder, generator, gmvaep, kl_warmup_callback, mmd_warmup_callback
+        return (
+            encoder,
+            generator,
+            grouper,
+            gmvaep,
+            kl_warmup_callback,
+            mmd_warmup_callback,
+        )
 
 
 # TODO: