<jupyter_start><jupyter_text>DatasetMNIST handwritten digits, reshaped to vectors in $\mathbb{R}^{784}$, and binarized.<jupyter_code>(x_train, _), (x_test, _) = tf.keras.datasets.mnist.load_data()

x_train = x_train.reshape((x_train.shape[0], -1)) / 255
x_train = np.where(x_train < 0.5, 0, 1).astype('float32')
x_test  = x_test.reshape((x_test.shape[0], -1)) / 255
x_test = np.where(x_test < 0.5, 0, 1).astype('float32')<jupyter_output><empty_output><jupyter_text>Variational Autoencoder EncoderThe encoder takes an input $x$, and compresses it down to a lower dimensonal latent vector $z$. Actually, it is stochastic so it outputs parameters $\theta = \{\mu, \sigma\}$ to $q_{\theta}(z|x)$ which is an isotropic multivariate Gaussian distribution $\mathcal{N}(\mu, \sigma I)$. When we sample from this distribution we get noisy representations of $z$. DecoderThe decoder takes a latent vector $z$ and outputs parameters for the data distributions that we can sample from to generate a sample. In this case of binarized, flattened MNIST images, the output is 784 Bernoulli parameters $p_{\phi}(x|z) = \{ \mu_1, \dots, \mu_{784} \}$. Variational AutoencoderThe full variational autoencoder has an encoder and a decoder. It takes a batch of input images in and outputs a batch of parameters $\mu, \sigma, p$ to be used to calculate the loss. $z$ is sampled using the reparameterization trick by letting $z=\mu + \sigma \epsilon$, where $\epsilon \sim \mathcal{N}(0, 1)$, to make $z$ deterministic with respect to the parameters, allowing gradients to be computed. Since the KL-diveregence regularizer pushes the parameters of the encoder toward a unit gaussian, we can then hallucinate MNIST digits by sampling from a unit gaussian (of the hidden dimension) and feeding these vectors through the decoder!<jupyter_code>class Encoder(tf.keras.layers.Layer):
    """ Gaussian distribution q(z|x) of dimension hidden_dim """
    def __init__(self, hidden_dim):
        super().__init__()
        self.h = hidden_dim
        self.dense1 = tf.keras.layers.Dense(2 * hidden_dim, activation='relu')
        self.dense2 = tf.keras.layers.Dense(2 * hidden_dim)
    
    def call(self, x):
        x = self.dense1(x)
        x = self.dense2(x)
        mean, std = x[:, :self.h], tf.nn.softplus(x[:, self.h:])
        return mean, std
        
    
class Decoder(tf.keras.layers.Layer):
    """ Bernoulli distribution p(x|z) of dimension output_dim """
    def __init__(self, hidden_dim, output_dim=784):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(2 * hidden_dim, activation='relu')
        self.dense2 = tf.keras.layers.Dense(output_dim, activation='sigmoid')
    
    def call(self, z):
        z = self.dense1(z)
        p = self.dense2(z)
        return p
        
    
class VAE(tf.keras.Model):
    
    def __init__(self, hidden_dim):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.encoder = Encoder(hidden_dim)
        self.decoder = Decoder(hidden_dim)
        
    def call(self, x):
        mean, std = self.encoder(x)
        eps = tf.random.normal(shape=mean.shape)
        z = mean + std * eps  # Reparametrization trick
        p = self.decoder(z)
        return mean, std, p
    
    def hallucinate(self, n):
        z = tf.random.normal(shape=(n, self.hidden_dim))
        p = self.decoder(z)
        return p.numpy().reshape(n, 28, 28)<jupyter_output><empty_output><jupyter_text>Loss FunctionImplementing the loss functions, the reconstruction loss is a single sample Monte-Carlo estimate and the KL-divergence loss is the analytical solution between two Gaussians. One is the unit Gaussian, and the other is the distribution output by the encoder.<jupyter_code>def log(x):
    """ Dodging nans """
    eps = 1e-8
    return tf.math.log(x + eps)


def reconstruction_loss(x, p):
    return -tf.einsum("ij, ij -> i", x, log(p)) - tf.einsum("ij, ij ->i", (1 - x), log(1 - p))


def kldiv_loss(mean, std):
    return tf.reduce_sum((tf.square(std) + tf.square(mean) - 1) / 2 - log(std), axis=-1)


def compute_loss(model, x):
    mean, std, p = model(x)
    rec_loss = reconstruction_loss(x, p)
    kld_loss = kldiv_loss(mean, std)
    return tf.reduce_mean(rec_loss + kld_loss)<jupyter_output><empty_output><jupyter_text>Helper FunctionsJust some helper functions to make the training loop a bit more readable!<jupyter_code>train_loss = tf.keras.metrics.Mean(name='training loss')
test_loss  = tf.keras.metrics.Mean(name='test loss')

def train_step(model, x):
    with tf.GradientTape() as tape:
            loss = compute_loss(model, x)
    
    grads = tape.gradient(loss, model.trainable_variables)
    opt.apply_gradients(zip(grads, model.trainable_variables))

    train_loss(loss)

    
def test_step(model, x):
    loss = compute_loss(model, x)
    test_loss(loss)

    
def reset_metrics():
    train_loss.reset_states()
    test_loss.reset_states()

    
def write_to_tensorboard(epoch):
    tf.summary.scalar(name='training loss', data=train_loss.result(), step=epoch)
    tf.summary.scalar(name='test loss', data=test_loss.result(), step=epoch)

    
def hallucinate_images(model, N=5, save=False):
    """ Show N * N hallucinated images in a grid """
    fig = plt.figure(figsize=(10, 10))
    grid = ImageGrid(fig, 111, nrows_ncols=(N, N), axes_pad=0.1)
    images = model.hallucinate(int(N * N))

    for ax, im in zip(grid, images):
        ax.imshow(im)
        ax.set_yticklabels([])
        ax.set_xticklabels([])
        ax.grid(False)

    plt.axis('off')
    if save:
        plt.savefig("hallucinated.png", dpi=500)
    plt.show()<jupyter_output><empty_output><jupyter_text>Training Loop<jupyter_code>EPOCHS = 25
LATENT = 50

bs = 128
test_ds = tf.data.Dataset.from_tensor_slices(x_test).batch(bs)
train_ds = tf.data.Dataset.from_tensor_slices(x_train).shuffle(10000).batch(bs)

current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
log_dir = 'logs/VAE/' + current_time
writer = tf.summary.create_file_writer(log_dir)

tmp_train = "[{:3d}, {:3d}] ---> Train loss: {:3.2f}"
tmp_test = "Test loss : {:3.2f}\n"

model = VAE(hidden_dim=LATENT)
opt = tf.keras.optimizers.Adam(1e-4)

reset_metrics()
with writer.as_default():
    
    for epoch in range(EPOCHS):

        for i, x_tr in enumerate(train_ds):
            train_step(model, x_tr)
            if i % 100 == 0:
                print(tmp_train.format(epoch, i, train_loss.result()))

        for x_te in test_ds:
            test_step(model, x_te)
        
        print(f"After epoch {epoch}:")
        print(tmp_test.format(test_loss.result()))
        
        # Visualize some hallucinated images every 10th epoch to see the progress!
        if epoch % 10 == 0:
            hallucinate_images(model, 5)
            
        write_to_tensorboard(epoch)
        reset_metrics()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12638483965014577, 'fraction_numerical': 0.014868804664723033, 'mean_word_length': 3.398076923076923, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '447404', 'license': 'MIT', 'n_tokens_mistral': 2423, 'n_tokens_neox': 2294, 'n_words': 601, 'path': 'VAE.ipynb', 'repo_name': 'AlexHermansson/variational-autoencoder', 'size': '6860'}