cleaning code and comments

bin/train/train_ed2esp.py

@@ -1,9 +1,7 @@
 ##########################################################################################
 #
-# This script is a sort of main file for the model in src.models.VAEAttention.py, which is 
-# a variational autoencoder using self-attention. This script will aim to show to train 
-# it, and also to save and load the model.
-# This file is modified to train in Dragonsoop
+# This script is a sort of main file for ED2ESP (FCN) model. This script will aim to 
+# show to train it, and also to save and load the model.
 #
 # Author: Juan Manuel Parrilla Gutierrez ([email protected])
 #
@@ -14,7 +12,7 @@
 import tensorflow as tf
 
 from src.utils.TFRecordLoader import TFRecordLoader
-from src.models.VAE_ed_esp import VAE_ed_esp
+from src.models.ED2ESP import VAE_ed_esp
 
 # RUN PARAMS #############################################################################
 os.environ["CUDA_VISIBLE_DEVICES"] = '2,3,4,5'

bin/train/train_vae.py

@@ -1,8 +1,7 @@
 ##########################################################################################
 #
-# This script is a sort of main file for the model in src.models.VAEAttention.py, which is 
-# a variational autoencoder using self-attention. This script will aim to show to train 
-# it, and also to save and load the model.
+# This script is a sort of main file for the model in src.models.VAE.py, and other VAE
+# variants. This script will aim to show to train it, and also to save and load the model.
 # This file is modified to train in Dragonsoop
 #
 # Author: Juan Manuel Parrilla Gutierrez ([email protected])

src/models/VAE_ed_esp.py → src/models/ED2ESP.py

@@ -58,19 +58,13 @@ def losses(self, data):
 
         x = self.preprocess_data(data[0])
         y = self.preprocess_esp(data[1])
-
-        # z_mean, z_log_var, z = self.encoder(x)
+
         z = self.encoder(x)
         y_nn = self.decoder(z)
         reconstruction_loss = tf.reduce_mean(
             tf.square(y - y_nn), axis=[1, 2, 3, 4]
         )
-        # reconstruction_loss *= self.r_loss_factor
-        # kl_loss = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
-        # kl_loss = tf.reduce_sum(kl_loss, axis=1)
-        # kl_loss *= -0.5
-        # total_loss = reconstruction_loss + kl_loss
-        # return total_loss, reconstruction_loss, kl_loss
+
         return reconstruction_loss
 
     def train_step(self, data):
@@ -130,27 +124,12 @@ def _build(self):
             x = conv_block(x, kernel_size, filters, stage=i, block='a', strides=strides)
             x = identity_block(x, kernel_size, filters, stage=i, block='b')
 
-        # shape_before_flattening = K.int_shape(x)[1:]
-
-        # x = Flatten()(x)
-        # self.mu = Dense(self.z_dim, name='mu')(x)
-        # self.log_var = Dense(self.z_dim, name='log_var')(x)
-
-        # self.z = Sampling(name='encoder_output')([self.mu, self.log_var])
-
-        # self.encoder = Model(
-        #     encoder_input, [self.mu, self.log_var, self.z], name='encoder'
-        #     )
-        self.encoder = Model (encoder_input, x, name='encoder')
+        self.encoder = Model(encoder_input, x, name='encoder')
 
         # THE DECODER
-        # decoder_input = Input(shape=(self.z_dim,), name='decoder_input')
         decoder_input = Input(shape=K.int_shape(x)[1:], name='decoder_input')
         x = decoder_input
 
-        # x = Dense(np.prod(shape_before_flattening))(decoder_input)
-        # x = Reshape(shape_before_flattening)(x)
-
         for i in range(self.n_layers_decoder):
             # just fetch the parameters in a variable so it doesn't get super long
             filters = self.decoder_conv_t_filters[i]

src/models/VAEattention.py

@@ -54,7 +54,7 @@ def _build(self):
             x = conv_block(x, kernel_size, fmaps, stage=i, block='a', strides=strides)
             x = identity_block(x, kernel_size, fmaps, stage=i, block='b')
 
-            # last layer, add residual blocks
+            # last layer, add transformer block
             if i == (self.n_layers_encoder-1):
                 x = TransformerBlock(filters)(x)
                 # x = TransformerBlock(filters)(x)