Update model.py

model.py

@@ -0,0 +1,153 @@
+import tensorflow as tf
+
+
+CAP = 3501 # Cap for the number of notes
+
+class Encoder_Z(tf.keras.layers.Layer):
+
+  def __init__(self, dim_z, name="encoder", **kwargs):
+    super(Encoder_Z, self).__init__(name=name, **kwargs)
+    self.dim_x = (3, CAP, 1)
+    self.dim_z = dim_z
+
+  def build(self):
+    layers = [tf.keras.layers.InputLayer(input_shape=self.dim_x)]
+
+    layers.append(tf.keras.layers.Conv2D(filters=64, kernel_size=3, strides=(2, 2)))
+    layers.append(tf.keras.layers.ReLU())
+    layers.append(tf.keras.layers.Flatten())
+
+    layers.append(tf.keras.layers.Dense(2000))
+    layers.append(tf.keras.layers.ReLU())
+
+    layers.append(tf.keras.layers.Dense(500))
+    layers.append(tf.keras.layers.ReLU())
+
+    layers.append(tf.keras.layers.Dense(self.dim_z * 2, activation=None, name="dist_params"))
+
+    return tf.keras.Sequential(layers)
+
+
+class Decoder_X(tf.keras.layers.Layer):
+
+  def __init__(self, dim_z, name="decoder", **kwargs):
+    super(Decoder_X, self).__init__(name=name, **kwargs)
+    self.dim_z = dim_z
+
+  def build(self):
+    # Build architecture
+      
+    layers = [tf.keras.layers.InputLayer(input_shape=(self.dim_z,))]
+
+    layers.append(tf.keras.layers.Dense(500))
+    layers.append(tf.keras.layers.ReLU())
+
+    layers.append(tf.keras.layers.Dense(2000))
+    layers.append(tf.keras.layers.ReLU())
+
+    layers.append(tf.keras.layers.Dense((CAP - 1) / 2 * 32, activation=None))
+    layers.append(tf.keras.layers.Reshape((1, int((CAP - 1) / 2), 32)))
+
+    layers.append(tf.keras.layers.Conv2DTranspose(
+        filters=64, kernel_size=3, strides=2, padding='valid'))
+    layers.append(tf.keras.layers.ReLU())
+
+    layers.append(tf.keras.layers.Conv2DTranspose(
+        filters=1, kernel_size=3, strides=1, padding='same'))
+
+    return tf.keras.Sequential(layers)
+
+kl_weight = tf.keras.backend.variable(0.125)
+
+
+
+class VAECost:
+    # VAE cost with a schedule based on the Microsoft Research Blog's article
+    # "Less pain, more gain: A simple method for VAE training with less of that KL-vanishing agony"
+    #
+    # The KL weight increases linearly, until it meets a certain threshold and keeps constant
+    # for the same number of epochs. After that, it decreases abruptly to zero again, and the
+    # cycle repeats.
+
+  def __init__(self, model):
+    self.model = model
+    self.kl_weight_increasing = True
+    self.epoch = 1
+
+
+  # The loss should have the form loss(y_true, y_pred), but in this
+  # case y_pred is computed in the cost function
+
+  @tf.function()
+  def __call__(self, x_true):
+    x_true = tf.cast(x_true, tf.float32)
+      
+    # Encode "song map" to get its latent representation and the parameters
+    # of the distribution
+    z_sample, mu, sd = self.model.encode(x_true)
+
+    # Decode the latent representation. Due to the VAE architecture, we should
+    # ideally get a reconstructed song map similar to the input.
+    x_recons = self.model.decoder(z_sample)
+
+    # Compute mean squared error, where our ground truth is the song map
+    # we pass as input, so we "compare" the reconstruction to it.
+      
+    recons_error = tf.cast(
+        tf.reduce_mean((x_true - x_recons) ** 2, axis=[1, 2, 3]),
+        tf.float32)
+
+    # Compute reverse KL divergence
+    kl_divergence = -0.5 * tf.math.reduce_sum(
+          1 + tf.math.log(tf.math.square(sd)) - tf.math.square(mu) - tf.math.square(sd),
+          axis=1) # shape=(batch_size,)
+
+    # Return metrics
+    elbo = tf.reduce_mean(-kl_weight * kl_divergence - recons_error)
+    mean_kl_divergence = tf.reduce_mean(kl_divergence)
+    mean_recons_error = tf.reduce_mean(recons_error)
+
+    return -elbo, mean_kl_divergence, mean_recons_error
+
+
+class VAE(tf.keras.Model):
+
+  def __init__(self, dim_z, seed=2000, analytic_kl=True, name="autoencoder", **kwargs):
+    super(VAE, self).__init__(name=name, **kwargs)
+    self.dim_x = (3, CAP, 1)
+    self.dim_z = dim_z
+    self.seed = seed
+    self.analytic_kl = analytic_kl
+    self.encoder = Encoder_Z(dim_z=self.dim_z).build()
+    self.decoder = Decoder_X(dim_z=self.dim_z).build()
+    self.cost_func = VAECost(self)
+
+  @tf.function()
+  def train_step(self, data):
+    # Gradient descent
+      
+    with tf.GradientTape() as tape:
+      neg_elbo, mean_kl_divergence, mean_recons_error = self.cost_func(data)
+
+    gradients = tape.gradient(neg_elbo, self.trainable_variables)
+    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
+
+    return {"abs ELBO": neg_elbo, "mean KL": mean_kl_divergence,
+            "mean recons": mean_recons_error,
+            "kl weight": kl_weight}
+
+  def encode(self, x_input):
+    # Get a "song map" and make a forward pass through the encoder, in order
+    # to return the latent representation and the distribution's parameters
+      
+    mu, rho = tf.split(self.encoder(x_input), num_or_size_splits=2, axis=1)
+    sd = tf.math.log(1 + tf.math.exp(rho))
+    z_sample = mu + sd * tf.random.normal(shape=(self.dim_z,))
+    return z_sample, mu, sd
+
+  def generate(self, z_sample=None):
+    # Decode a latent representation of a song, which is provided or sampled
+      
+    if z_sample == None:
+      z_sample = tf.expand_dims(tf.random.normal(shape=(self.dim_z,)), axis=0)
+    return self.decoder(z_sample)