File size: 5,510 Bytes
92ac48c 5a0bd26 92ac48c 523a819 92ac48c 523a819 92ac48c 523a819 92ac48c 661ebc2 92ac48c 0526e94 5a0bd26 92ac48c 5a0bd26 92ac48c 5a0bd26 92ac48c 5a0bd26 92ac48c |
1 2 3 4 5 6 7 8 9 10 11 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 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 |
import tensorflow as tf
import os
import inspect
_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, name="variational autoencoder", **kwargs):
super(VAE, self).__init__(name=name, **kwargs)
self.dim_x = (3, _CAP, 1)
self.encoder = Encoder_Z(dim_z=120).build()
self.decoder = Decoder_X(dim_z=120).build()
self.cost_func = VAECost(self)
# Get the path of the script that defines this method
script_path = inspect.getfile(inspect.currentframe())
# Get the directory containing the script
script_dir = os.path.dirname(os.path.abspath(script_path))
# Construct the path to the weights folder
weights_dir = os.path.join(script_dir, 'weights') + os.sep
# Load pretrained weights
self.load_weights(weights_dir)
@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=(120,))
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=(120,)), axis=0)
return self.decoder(z_sample)
|