More inputs

model/encoder.py

@@ -5,7 +5,8 @@
 
 # Third party module imports
 from keras import Input, Model
-from keras.layers import BatchNormalization, Dense, Dropout, GlobalAveragePooling1D
+from keras.layers import BatchNormalization, Dense, Dropout, GlobalAveragePooling1D, \
+    LayerNormalization
 
 # Local module imports
 from model.transformer import TimeseriesTransformerBuilder as TSTFBuilder
@@ -87,6 +88,7 @@ def _modelstack(self, input_shape, head_size, num_heads, ff_dim,
         for dim in mlp_units:
             x = Dense(dim, activation="relu")(x)
             x = Dropout(mlp_dropout)(x)
+            x = LayerNormalization()(x)
 
         # Two separate latent spaces supported
 

model/latent.py

@@ -81,7 +81,7 @@ def _build(self, mlp_dim, semantic_dims):
         ]
 
         # Compute MAE, with normalization by approximate range of values (~4)
-        errors.append(mean_absolute_error(targets[-1], reg) / 4.)
+        errors.append(mean_absolute_error(targets[-1], reg) / 40.)
         error = ops.sum(ops.stack(errors))
 
         # Combine everything into single dense layer

model/model.py

@@ -50,7 +50,7 @@ def __init__(self, input_shape, head_size, num_heads, ff_dim,
         self.cycle_len = 25
         self.total_epoch = None
         self.max_cap = 5.0
-        self.beta = 0.01
+        self.beta = 0.05
 
         self.encoder = Encoder(
                 input_shape,
@@ -146,7 +146,7 @@ def kl_capacity(self, epoch, max_cap, total_epochs):
         cap = max_cap * epoch / total_epochs
         return ops.minimum(max_cap, cap)
 
-    def kl_weight(self, epoch, step, total_epochs, cycle, max_cap):
+    def kl_weight(self, epoch, total_epochs, cycle, max_cap):
         beta = (
             self.beta_anneal(epoch, total_epochs) *
             self.beta_cyclical(epoch, cycle)
@@ -156,36 +156,29 @@ def kl_weight(self, epoch, step, total_epochs, cycle, max_cap):
 
     """
     def train_step(self, data, sample_weight=None):
-
-        with ops.GradientTape() as tape:
-            outputs, sem_loss, kl_loss, recon_loss = self(
-                    data, training=True
-            )
-            self.recon_weight = ops.cast(1./4., dtype=ops.float32)
-            epoch = ops.cast(
-                    self.optimizer.iterations // self.steps_per_epoch,
-                    kl_loss.dtype
-            )
-            step = ops.cast(self.optimizer.iterations, kl_loss.dtype)
-            beta, capacity = self.kl_weight(
-                    epoch=epoch,
-                    step=step,
-                    total_epochs=self.total_epoch,
-                    cycle=self.cycle_len,
-                    max_cap=self.max_cap
-            )
-            kl_term = beta * ops.abs(kl_loss - capacity)
-            total_loss = sem_loss + kl_term + recon_loss * self.recon_weight
+        outputs, sem_loss, kl_loss, recon_loss = self(data)
+        beta, capacity = self.kl_weight(
+                epoch=self.current_epoch,
+                total_epochs=self.total_epoch,
+                cycle=self.cycle_len,
+                max_cap=self.max_cap
+        )
+        
+        kl_weighted = beta * ops.abs(kl_loss - capacity)
+        total_loss = ops.mean(
+                sem_loss + kl_weighted + recon_loss * self.recon_weight
+        )
 
         trainable_vars = self.trainable_variables
-        grads = tape.gradient(total_loss, trainable_vars)
-        self.optimizer.apply_gradients(zip(grads, trainable_vars))
+        grads = ops.gradient(total_loss, trainable_vars)
+        self.optimizer.apply(grads, trainable_vars)
+        self.compiled_metrics.update_state(data[0], outputs)
 
         return {
             "loss": total_loss,
-            "sem_loss": sem_loss,
-            "kl_loss": kl_loss,
-            "recon_loss": recon_loss,
+            "sem_loss": ops.mean(sem_loss),
+            "kl_loss": ops.mean(kl_loss),
+            "recon_loss": ops.mean(recon_loss),
             "beta": beta,
             "capacity": capacity
         }"""