assignment11_starter.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Imports, define AE model\"\"\"\n",
    "import tensorflow as tf\n",
    "tfkl = tf.keras.layers\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import image as mpimage\n",
    "\n",
    "from data.utils import parse_image_example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# just a skeleton.\n",
    "# you fill in the autoencoding stuff ;)\n",
    "# this is mainly the VQ-related things\n",
    "# if you find it awkard to work with half-finished code written by someone else,\n",
    "# feel free to roll your own version.\n",
    "class VQVAE(tf.keras.Model):\n",
    "    def __init__(self, ..., **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        #...\n",
    "        \n",
    "        # this is to track unused codebook vectros\n",
    "        self.usage_tracker = tf.keras.metrics.MeanTensor(\"codebook_usage\")\n",
    "        \n",
    "        dim_code = encoder.output_shape[-1]\n",
    "        self.codebook = tf.Variable(tf.random.normal([codebook_size, dim_code]))\n",
    "        self.codebook_size = codebook_size\n",
    "        \n",
    "    def call(self, inputs, training=None):\n",
    "        return self.apply_ae(inputs, training=training)[0]\n",
    "    \n",
    "    def train_step(self, data):\n",
    "        with tf.GradientTape() as tape:\n",
    "            reconstruction, encoder_outputs, codes, indices = self.apply_ae(data, training=True)\n",
    "            \n",
    "            reconstruction_loss = ...\n",
    "            codebook_loss = tf.reduce_mean(tf.square(tf.stop_gradient(encoder_outputs) - codes))*code_shape[-1]\n",
    "            commitment_loss = tf.reduce_mean(tf.square(encoder_outputs - tf.stop_gradient(codes)))*code_shape[-1]\n",
    "            total_loss = reconstruction_loss + codebook_loss + self.beta*commitment_loss\n",
    "            \n",
    "        variables = self.encoder.trainable_variables + self.decoder.trainable_variables + [self.codebook]\n",
    "        gradients = tape.gradient(total_loss, variables)\n",
    "        self.optimizer.apply_gradients(zip(gradients, variables))\n",
    "        \n",
    "        # track codebook usage\n",
    "        flat_indices = tf.one_hot(tf.reshape(indices, (-1,)), depth=self.codebook_size)\n",
    "        usage_count = tf.reduce_sum(flat_indices, axis=0)\n",
    "    \n",
    "        self.usage_tracker.update_state(usage_count)\n",
    "        \n",
    "        # update with other metrics, i.e. losses\n",
    "        # see e.g. https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit\n",
    "        return {\"codebook_usage\": self.usage_tracker.result()}\n",
    "    \n",
    "    def test_step(self, data):\n",
    "        reconstruction, encoder_outputs, codes, indices = self.apply_ae(data, training=False)\n",
    "        \n",
    "        reconstruction_loss = ...\n",
    "        codebook_loss = tf.reduce_mean(tf.square(tf.stop_gradient(encoder_outputs) - codes))*code_shape[-1]\n",
    "        commitment_loss = tf.reduce_mean(tf.square(encoder_outputs - tf.stop_gradient(codes)))*code_shape[-1]\n",
    "        total_loss = reconstruction_loss + codebook_loss + self.beta*commitment_loss\n",
    "        \n",
    "        # track codebook usage\n",
    "        flat_indices = tf.one_hot(tf.reshape(indices, (-1,)), depth=self.codebook_size)\n",
    "        usage_count = tf.reduce_sum(flat_indices, axis=0)\n",
    "    \n",
    "        self.usage_tracker.update_state(usage_count)\n",
    "        \n",
    "        return {\"codebook_usage\": self.usage_tracker.result()}\n",
    "    \n",
    "    @property\n",
    "    def metrics(self):\n",
    "        return [self.usage_tracker]\n",
    "    \n",
    "    # vq-vae-specific stuff\n",
    "    def read_codebook(self, encoder_outputs, with_indices=False):\n",
    "        distances = tf.reduce_mean(tf.square(encoder_outputs[:, :, :, None, :] \n",
    "                                             - self.codebook[None, None, None, :, :]), axis=-1)\n",
    "        #print(\"d\", distances.shape)\n",
    "        min_distance_inds = tf.math.argmin(distances, axis=-1)\n",
    "        #print(\"ind\", min_inds.shape)\n",
    "        codes = tf.gather(self.codebook, min_distance_inds)\n",
    "\n",
    "        # 1st output are the codebooks with \"straight-through estimator\".\n",
    "        # this allows gradients to flow from the decoder into the encoder.\n",
    "        # to learn the codebook, we also return codes, as these allow gradient flow into the codebook.\n",
    "        if with_indices:\n",
    "            return encoder_outputs + tf.stop_gradient(codes - encoder_outputs), codes, min_distance_inds\n",
    "        else:\n",
    "            return encoder_outputs + tf.stop_gradient(codes - encoder_outputs), codes\n",
    "\n",
    "    # reconstructions for reconstruction loss (duh)\n",
    "    # encoder outputs and codes for codebook and commitment loss\n",
    "    # indices because we may need them to check usage, create training data for AR model etc\n",
    "    def apply_ae(self, inputs, training=False):\n",
    "        encoder_outputs = self.encoder(inputs, training=training)\n",
    "        codes, codes_with_gradients, indices = self.read_codebook(encoder_outputs, with_indices=True)\n",
    "        return (decoder(codes, training=training),\n",
    "                encoder_outputs,\n",
    "                codes_with_gradients,\n",
    "                indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = ...  # model\n",
    "decoder = ... # another model\n",
    "vqvae = VQVAE(encoder, decoder, codebook_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# this can be used to reset unused codebook entries after each epoch\n",
    "# ...or less frequently (frequency argument)\n",
    "class CodebookResetter(tf.keras.callbacks.Callback):\n",
    "    def __init__(self, frequency, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.frequency = frequency\n",
    "        self.reference_batch = next(iter(train_data))\n",
    "    \n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        if not epoch % self.frequency:\n",
    "            average_usage = logs[\"codebook_usage\"]\n",
    "            # entries that were never used this epoch\n",
    "            unused_code_indices = np.where(average_usage == 0)[0]\n",
    "            print(\"\\nDEBUG UNUSED INDICES\", unused_code_indices)\n",
    "            \n",
    "            # we take random encoder outputs for a reference batch\n",
    "            reference_encodings = self.model.encoder(self.reference_batch)\n",
    "            \n",
    "            new_codebook_entries = tf.stack([reference_encodings[np.random.choice(batch_size),\n",
    "                                               np.random.choice(code_image_w),\n",
    "                                               np.random.choice(code_image_h)] for _ in unused_code_indices],\n",
    "                                            axis=0)\n",
    "            print(\"DEBUG NEW ENTRIES\", new_codebook_entries.shape, \"\\n\")\n",
    "            \n",
    "            # unused entries are replaced by reference encodings\n",
    "            if len(unused_code_indices):\n",
    "                sparse_update = tf.IndexedSlices(new_codebook_entries,\n",
    "                                                 tf.convert_to_tensor(unused_code_indices, dtype=tf.int32))\n",
    "                self.model.codebook.scatter_update(sparse_update)\n",
    "\n",
    "\n",
    "# fit..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a dataset of codes\n",
    "# NOTE this stores both the code vectors AND the indices.\n",
    "# you may also store only the indices (takes less space)\n",
    "\n",
    "train_code_inds = []\n",
    "train_codes = []\n",
    "for step, img_batch in enumerate(train_data):\n",
    "    encodings = encoder(img_batch)\n",
    "    codes, _, inds = vqvae.read_codebook(encodings, with_indices=True)\n",
    "    train_code_inds.append(inds)\n",
    "    train_codes.append(codes)\n",
    "train_code_inds = np.concatenate(train_code_inds)\n",
    "train_codes = np.concatenate(train_codes)\n",
    "\n",
    "code_data = tf.data.Dataset.from_tensor_slices((train_code_inds, train_codes))\n",
    "\n",
    "code_data_train = code_data.shuffle(60000).batch(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a dataset of codes (test set)\n",
    "\n",
    "test_code_inds = []\n",
    "test_codes = []\n",
    "for step, img_batch in enumerate(test_data):\n",
    "    if not step % 50:\n",
    "        print(step)\n",
    "    encodings = encoder(img_batch)\n",
    "    codes, _, inds = vqvae.read_codebook(encodings, with_indices=True)\n",
    "    test_code_inds.append(inds)\n",
    "    test_codes.append(codes)\n",
    "test_code_inds = np.concatenate(test_code_inds)\n",
    "test_codes = np.concatenate(test_codes)\n",
    "\n",
    "code_data_test = tf.data.Dataset.from_tensor_slices((test_code_inds, test_codes)).batch(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bonus: check codebook usage. are all 256 vectors used, and how often?\n",
    "from collections import Counter\n",
    "\n",
    "usage_count = Counter(train_code_inds.reshape((-1,)))\n",
    "\n",
    "descending = [thing[1] for thing in usage_count.most_common()]\n",
    "plt.bar(range(len(descending)), descending)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now train a pixelcnn on the codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tested for TF 2.11!\n",
    "# it may not work in other version :(\n",
    "from tensorflow.python.keras.layers.convolutional import Conv\n",
    "\n",
    "class MaskedConv2D(tfkl.Conv2D):\n",
    "  def __init__(self,\n",
    "               filters,\n",
    "               kernel_size,\n",
    "               strides=1,\n",
    "               padding='valid',\n",
    "               data_format=None,\n",
    "               dilation_rate=1,\n",
    "               activation=None,\n",
    "               use_bias=True,\n",
    "               kernel_initializer='glorot_uniform',\n",
    "               bias_initializer='zeros',\n",
    "               kernel_regularizer=None,\n",
    "               bias_regularizer=None,\n",
    "               activity_regularizer=None,\n",
    "               kernel_constraint=None,\n",
    "               bias_constraint=None,\n",
    "               trainable=True,\n",
    "               name=None,\n",
    "               mask_type=\"a\",\n",
    "               **kwargs):\n",
    "    super(MaskedConv2D, self).__init__(\n",
    "               filters,\n",
    "               kernel_size,\n",
    "               strides=strides,\n",
    "               padding=padding,\n",
    "               data_format=data_format,\n",
    "               dilation_rate=dilation_rate,\n",
    "               activation=activation,\n",
    "               use_bias=use_bias,\n",
    "               kernel_initializer=kernel_initializer,\n",
    "               bias_initializer='zeros',\n",
    "               kernel_regularizer=None,\n",
    "               bias_regularizer=None,\n",
    "               activity_regularizer=None,\n",
    "               kernel_constraint=None,\n",
    "               bias_constraint=None,\n",
    "               trainable=True,\n",
    "               name=None,\n",
    "        **kwargs)\n",
    "    self.mask = np.zeros([kernel_size, kernel_size, 1, 1], dtype=np.float32)\n",
    "    self.mask[:kernel_size//2, :, :, :] = 1\n",
    "    self.mask[kernel_size//2, :kernel_size//2, :, : ] = 1\n",
    "    if mask_type == \"b\":\n",
    "        self.mask[kernel_size//2, kernel_size//2] = 1\n",
    "    self.mask = tf.convert_to_tensor(self.mask)\n",
    "    \n",
    "    \n",
    "  def call(self, inputs):\n",
    "    masked_kernel = self.mask * self.kernel\n",
    "    outputs = self.convolution_op(inputs, masked_kernel)\n",
    "\n",
    "    if self.use_bias:\n",
    "      if self.data_format == 'channels_first':\n",
    "        if self.rank == 1:\n",
    "          # nn.bias_add does not accept a 1D input tensor.\n",
    "          bias = tf.python.ops.array_ops.reshape(self.bias, (1, self.filters, 1))\n",
    "          outputs += bias\n",
    "        else:\n",
    "          outputs = tf.nn.bias_add(outputs, self.bias, data_format='NCHW')\n",
    "      else:\n",
    "        outputs = tf.nn.bias_add(outputs, self.bias, data_format='NHWC')\n",
    "\n",
    "    if self.activation is not None:\n",
    "      return self.activation(outputs)\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# proposal for a pixelcnn class.\n",
    "# training is catually quite basic. the masked convolutions are set up such that\n",
    "# the output at a given position is the prediction for that pixel.\n",
    "# since the 1st layer mask also removes the center pixel, this is fine.\n",
    "# we don't have to shift inputs vs targets by one step, like in an autoregressive RNN for example.\n",
    "\n",
    "# note that this models takes the code VECTORS as input, but predicts the indices.\n",
    "# it would also be possible to change this: take the indices as vectors.\n",
    "# in that case, the first model layer should probably be an embedding.\n",
    "# this effectively allows the AR model to learn its own representations for the indices,\n",
    "# instead of re-using the code vectors from the VQ-VAE.\n",
    "class PixelCNN(tf.keras.Model):\n",
    "    def __init__(self, inputs, outputs, **kwargs):\n",
    "        super().__init__(inputs, outputs, **kwargs)\n",
    "        \n",
    "        self.cross_entropy_tracker = tf.keras.metrics.Mean(\"cross_entropy\")\n",
    "        \n",
    "    def train_step(self, data):\n",
    "        indices, codes = data\n",
    "        with tf.GradientTape() as tape:\n",
    "            outputs = self(codes, training=True)\n",
    "            loss = self.compiled_loss(indices, outputs)\n",
    "\n",
    "        variables = self.trainable_variables\n",
    "        gradients = tape.gradient(loss, variables)\n",
    "        self.optimizer.apply_gradients(zip(gradients, variables))\n",
    "        \n",
    "        self.cross_entropy_tracker.update_state(loss)\n",
    "        return {\"cross_entropy\": self.cross_entropy_tracker.result()}\n",
    "    \n",
    "    def test_step(self, data):\n",
    "        indices, codes = data\n",
    "        outputs = self(codes, training=False)\n",
    "        loss = self.compiled_loss(indices, outputs)\n",
    "        \n",
    "        self.cross_entropy_tracker.update_state(loss)\n",
    "        return {\"cross_entropy\": self.cross_entropy_tracker.result()}\n",
    "    \n",
    "    # generation proceeds pixel by pixel\n",
    "    def generate(self, num_samples):\n",
    "        image = np.zeros([num_samples, code_image_w, code_image_h, dim_code], dtype=np.float32)\n",
    "        for row in range(code_image_w):\n",
    "            for col in range(code_image_h):\n",
    "                index_logits = self(image)[:, row, col, :]\n",
    "                index_sample = tfd.Categorical(logits=index_logits, dtype=tf.int32).sample()\n",
    "                image[:, row, col] = tf.gather(vqvae.codebook, index_sample)\n",
    "\n",
    "        return map_for_likelihood(decoder(image), likelihood)\n",
    "    \n",
    "    @property\n",
    "    def metrics(self):\n",
    "        return [self.cross_entropy_tracker]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp = tf.keras.Input(code_shape)\n",
    "x = inp\n",
    "x = MaskedConv2D(32, 3, padding=\"same\", mask_type=\"a\")(x)\n",
    "# you can stack more MaskedConv2D layers, build residual blocks with 1x1 convolutions, etc.\n",
    "# note that you can use mask_type=\"b\" for all layers except the first one.\n",
    "x = ...\n",
    "# output layer: predict codebook index\n",
    "x = tfkl.Conv2D(codebook_size, 1, padding=\"same\")(x)\n",
    " \n",
    "pixel_cnn = PixelCNN(inp, x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# indices are classes, so use cross-entropy as loss\n",
    "pixel_loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "class ImageGenCallback(tf.keras.callbacks.Callback):\n",
    "    def __init__(self, frequency, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.frequency = frequency\n",
    "        \n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        if not epoch % self.frequency:\n",
    "            generated_batch = self.model.generate(64)\n",
    "        \n",
    "            plt.figure(figsize=(15,15))\n",
    "            for ind, image in enumerate(generated_batch):\n",
    "                plt.subplot(8, 8, ind+1)\n",
    "                plt.imshow(image)\n",
    "                plt.axis(\"off\")\n",
    "            plt.suptitle(\"Random generations\")\n",
    "            plt.show()\n",
    "            \n",
    "image_gen_callback = ImageGenCallback(10)\n",
    "            \n",
    "pixel_cnn.fit(code_data_train, validation_data=code_data_test, epochs=n_epochs, callbacks=[image_gen_callback])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = pixel_cnn.generate(64).numpy()\n",
    "plt.figure(figsize=(15,15))\n",
    "for ind, img in enumerate(samples):\n",
    "    plt.subplot(8, 8, ind+1)\n",
    "    plt.imshow(img)\n",
    "    plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}