discrete-VAE-experiments-neurips-identity.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# choose a particular GPU\n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow.keras.layers import Dense, Flatten\n",
    "from tensorflow.keras import Model\n",
    "from tensorflow.keras.utils import plot_model\n",
    "\n",
    "import tensorflow_probability as tfp\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class IdentitySubsetLayer(tf.keras.layers.Layer):\n",
    "  \n",
    "    def __init__(self, _k=10, _tau=10.0, _lambda=10.0):\n",
    "        super(IdentitySubsetLayer, self).__init__()\n",
    "        \n",
    "        self.k = _k\n",
    "        self._tau = _tau\n",
    "        self._lambda = _lambda\n",
    "        self.samples = None\n",
    "        self.gumbel_dist = tfp.distributions.Gumbel(loc=0.0, scale=1.0)\n",
    "        \n",
    "    @tf.function\n",
    "    def sample_gumbel(self, shape, eps=1e-20):\n",
    "        return self.gumbel_dist.sample(shape)\n",
    "    \n",
    "    @tf.function\n",
    "    def sample_gumbel_k(self, shape):\n",
    "        \n",
    "        s = tf.map_fn(fn=lambda t: tf.random.gamma(shape, 1.0/self.k,  self.k/t), \n",
    "                  elems=tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]))\n",
    "        # now add the samples\n",
    "        s = tf.reduce_sum(s, 0)\n",
    "        # the log(m) term\n",
    "        s = s - tf.math.log(10.0)\n",
    "        # divide by k --> each s[c] has k samples whose sum is distributed as Gumbel(0, 1)\n",
    "        s = self._tau * (s / self.k)\n",
    "\n",
    "        return s\n",
    "       \n",
    "    @tf.custom_gradient\n",
    "    def identity_layer(self, logits, hard=False):\n",
    "\n",
    "        # toggle between Gumbel(0, 1) and sum-of-Gamma\n",
    "        #forward_sample = self.sample_gumbel(tf.shape(logits))\n",
    "        threshold = tf.expand_dims(tf.nn.top_k(logits, self.k, sorted=True)[0][:,-1], -1)\n",
    "        y_map = tf.cast(tf.greater_equal(logits, threshold), tf.float32)\n",
    "\n",
    "        def custom_grad(dy):\n",
    "            return dy, hard\n",
    "\n",
    "        return y_map, custom_grad\n",
    "\n",
    "    def call(self, logits, hard=False):\n",
    "        logits = logits + self.sample_gumbel_k(tf.shape(logits))\n",
    "        # Project onto the hyperplane\n",
    "        logits = logits - tf.reduce_mean(logits, 1)[:, None]\n",
    "        # Project onto the sphere\n",
    "        logits = logits / tf.norm(logits, axis=1)[:, None]\n",
    "        return self.identity_layer(logits, hard)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "PARAMS = {\n",
    "    \"batch_size\": 100,\n",
    "    \"data_dim\": 784,\n",
    "    \"M\": 20,\n",
    "    \"N\": 20,\n",
    "    \"nb_epoch\": 100, \n",
    "    \"epsilon_std\": 0.01,\n",
    "    \"anneal_rate\": 0.0003,\n",
    "    \"init_temperature\": 1.0,\n",
    "    \"min_temperature\": 0.5,\n",
    "    \"learning_rate\": 1e-3,\n",
    "    \"hard\": False,\n",
    "}\n",
    "\n",
    "class DiscreteVAE(tf.keras.Model):\n",
    "    \n",
    "    def __init__(self, params):\n",
    "        super(DiscreteVAE, self).__init__()\n",
    "        \n",
    "        self.params = params\n",
    "                \n",
    "        # encoder\n",
    "        self.enc_dense1 = tf.keras.layers.Dense(512, activation='relu')\n",
    "        self.enc_dense2 = tf.keras.layers.Dense(256, activation='relu')\n",
    "        self.enc_dense3 = tf.keras.layers.Dense(params[\"N\"]*params[\"M\"])\n",
    "        \n",
    "        # this is our new Gumbel layer\n",
    "        self.identityLayer = IdentitySubsetLayer(_k=10, _tau=10.0, _lambda=10.0)\n",
    "\n",
    "        # decoder\n",
    "        self.flatten = Flatten()\n",
    "        self.dec_dense1 = tf.keras.layers.Dense(256, activation='relu')\n",
    "        self.dec_dense2 = tf.keras.layers.Dense(512, activation='relu')\n",
    "        self.dec_dense3 = tf.keras.layers.Dense(params[\"data_dim\"])\n",
    "\n",
    "\n",
    "    def sample_gumbel(self, shape, eps=1e-20): \n",
    "        \"\"\"Sample from Gumbel(0, 1)\"\"\" \n",
    "        # test\n",
    "        #U = gumbel_dist.sample(shape)\n",
    "        U = tf.random.uniform(shape, minval=0, maxval=1)\n",
    "        return -tf.math.log(-tf.math.log(U + eps) + eps)\n",
    "    \n",
    "    def gumbel_softmax_sample(self, logits, temperature): \n",
    "        \"\"\" Draw a sample from the Gumbel-Softmax distribution\"\"\"\n",
    "        # logits: [batch_size, n_class] unnormalized log-probs\n",
    "        y = logits + self.sample_gumbel(tf.shape(logits))\n",
    "        return tf.nn.softmax(y / temperature)  \n",
    "\n",
    "    def gumbel_softmax(self, logits, temperature, hard=True):\n",
    "        \"\"\"\n",
    "        logits: [batch_size, n_class] unnormalized log-probs\n",
    "        temperature: non-negative scalar\n",
    "        hard: if True, take argmax, but differentiate w.r.t. soft sample y\n",
    "        \"\"\"\n",
    "        y = self.gumbel_softmax_sample(logits, temperature)\n",
    "        if hard: \n",
    "            # \n",
    "            y_hard = tf.cast(tf.equal(y, tf.reduce_max(y, 1, keepdims=True)),y.dtype)\n",
    "            y = tf.stop_gradient(y_hard - y) + y\n",
    "        return y\n",
    "    \n",
    "    def decoder(self, x):\n",
    "        # decoder\n",
    "        h = self.flatten(x)\n",
    "        h = self.dec_dense1(h)\n",
    "        h = self.dec_dense2(h)\n",
    "        h = self.dec_dense3(h)\n",
    "        return h\n",
    "\n",
    "    def call(self, x, tau, hard=False):\n",
    "        N = self.params[\"N\"]\n",
    "        M = self.params[\"M\"]\n",
    "\n",
    "        # encoder\n",
    "        x = self.enc_dense1(x)\n",
    "        x = self.enc_dense2(x)\n",
    "        x = self.enc_dense3(x)   # (batch, N*M)\n",
    "        logits_y = tf.reshape(x, [-1, M])   # (batch*N, M)\n",
    "\n",
    "        ###################################################################\n",
    "        ## here we toggle between methods #################################\n",
    "        # here we can switch between traditional and our method\n",
    "        # \"traditional\" Gumbel Softmax trick\n",
    "        #y = self.gumbel_softmax(logits=logits_y, temperature=tau, hard=False)\n",
    "        # IMLE approach -- note: we don't anneal so set temperature once at init\n",
    "        y = self.identityLayer(logits=logits_y, hard=False)\n",
    "        ###################################################################\n",
    "        \n",
    "        assert y.shape == (self.params[\"batch_size\"]*N, M)\n",
    "        y = tf.reshape(y, [-1, N, M])\n",
    "        self.sample_y = y\n",
    "\n",
    "        # decoder\n",
    "        logits_x = self.decoder(y)\n",
    "        return logits_y, logits_x\n",
    "\n",
    "\n",
    "def gumbel_loss(model, x, tau, hard=True):\n",
    "    M = 20\n",
    "    N = 20\n",
    "    data_dim = PARAMS['data_dim']\n",
    "    logits_y, logits_x = model(x, tau, hard)\n",
    "    \n",
    "    # cross-entropy\n",
    "    cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(labels=x, logits=logits_x)\n",
    "    cross_ent = tf.math.reduce_sum(cross_ent, 1)\n",
    "    cross_ent = tf.math.reduce_mean(cross_ent, 0)\n",
    "    \n",
    "    # KL loss\n",
    "    q_y = tf.nn.softmax(logits_y)   # (batshsize*N, M)  softmax\n",
    "    log_q_y = tf.math.log(q_y + 1e-20)   # (batshsize*N, M)  \n",
    "    kl_tmp = tf.reshape(q_y*(log_q_y-tf.math.log(1.0/M)), [-1,N,M])  # (batch_size,N,K)\n",
    "    KL = tf.math.reduce_sum(kl_tmp, [1, 2])    # shape=(batch_size, 1)\n",
    "        \n",
    "    KL_mean = tf.math.reduce_mean(KL)\n",
    "    #print(\"**\", cross_ent.numpy(), KL_mean.numpy())\n",
    "    return cross_ent + KL_mean\n",
    "\n",
    "\n",
    "def compute_gradients(model, x, tau, hard):\n",
    "    with tf.GradientTape() as tape:\n",
    "        loss = gumbel_loss(model, x, tau, hard)\n",
    "    return tape.gradient(loss, model.trainable_variables), loss\n",
    "\n",
    "\n",
    "def apply_gradients(optimizer, gradients, variables):\n",
    "    optimizer.apply_gradients(zip(gradients, variables))\n",
    "\n",
    "\n",
    "def get_learning_rate(step, init=PARAMS[\"learning_rate\"]):\n",
    "    return tf.convert_to_tensor(init * pow(0.95, (step / 1000.)), dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-02-27 14:15:41.500572: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz\n",
      "2023-02-27 14:15:41.531197: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "tf.random.set_seed(1234)\n",
    "\n",
    "model = DiscreteVAE(PARAMS)\n",
    "# model.build()\n",
    "# plot_model(model, to_file='model_plot.pdf', show_shapes=True, show_layer_names=True)\n",
    "learning_rate = tf.Variable(PARAMS[\"learning_rate\"], trainable=False, name=\"LR\")\n",
    "\n",
    "optimizer = tf.keras.optimizers.Adam()\n",
    "\n",
    "# data\n",
    "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
    "x_train = x_train.astype('float32') / 255.\n",
    "x_test = x_test.astype('float32') / 255.\n",
    "x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))\n",
    "x_test = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))\n",
    "\n",
    "TRAIN_BUF = 60000\n",
    "BATCH_SIZE = 100\n",
    "TEST_BUF = 10000\n",
    "\n",
    "train_dataset = tf.data.Dataset.from_tensor_slices(x_train).shuffle(TRAIN_BUF).batch(BATCH_SIZE)\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices(x_test).shuffle(TEST_BUF).batch(BATCH_SIZE)\n",
    "\n",
    "# temperature\n",
    "tau = PARAMS[\"init_temperature\"]\n",
    "anneal_rate = PARAMS[\"anneal_rate\"]\n",
    "min_temperature = PARAMS[\"min_temperature\"]\n",
    "\n",
    "results = []\n",
    "\n",
    "# Train\n",
    "for epoch in range(1, PARAMS[\"nb_epoch\"] + 1):\n",
    "    \n",
    "    # this is only needed for the standard Gumbel softmax trick\n",
    "    tau = np.maximum(tau * np.exp(-anneal_rate*epoch), min_temperature)\n",
    "\n",
    "    for train_x in train_dataset:\n",
    "        gradients, loss = compute_gradients(model, train_x, tau, hard=PARAMS[\"hard\"])\n",
    "        apply_gradients(optimizer, gradients, model.trainable_variables)\n",
    "\n",
    "    print(\"Epoch:\", epoch, \", TRAIN loss:\", loss.numpy(), \", Temperature:\", tau)\n",
    "\n",
    "    if epoch % 1 == 0:\n",
    "        losses = []\n",
    "        for test_x in test_dataset:\n",
    "            losses.append(gumbel_loss(model, test_x, tau, hard=True))\n",
    "        eval_loss = np.mean(losses)\n",
    "        results.append(eval_loss)\n",
    "        print(\"Eval Loss:\", eval_loss, \"\\n\")\n",
    "\n",
    "    if PARAMS['hard'] == True:\n",
    "        model.save_weights(\"model.h5\")\n",
    "    else:\n",
    "        model.save_weights(\"model_hard.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(x_train, _), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
    "x_train = x_train.astype('float32') / 255.\n",
    "x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))\n",
    "x_test = x_test.astype('float32') / 255.\n",
    "x_test = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))\n",
    "\n",
    "# Model Forward\n",
    "logits_y, logits_x1 = model(x_test[10:110], tau=1.0, hard=True)\n",
    "logits_y, logits_x2 = model(x_test[10:110], tau=1.0, hard=True)\n",
    "logits_y, logits_x3 = model(x_test[10:110], tau=1.0, hard=True)\n",
    "\n",
    "\n",
    "sample_y = model.sample_y.numpy()\n",
    "\n",
    "logits_x1 = tf.sigmoid(logits_x1).numpy()\n",
    "logits_x2 = tf.sigmoid(logits_x2).numpy()\n",
    "logits_x3 = tf.sigmoid(logits_x3).numpy()\n",
    "\n",
    "print(\"sample_y: \", sample_y.shape, \", logits_x.shape:\", logits_x1.shape)\n",
    "code = model.sample_y\n",
    "\n",
    "\n",
    "def save_plt(original_img, construct_img1, construct_img2, construct_img3, code):\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    for i in range(0, 25, 5):\n",
    "        # input img\n",
    "        plt.subplot(5, 5, i+1)\n",
    "        plt.imshow(original_img[i, :].reshape(28, 28), cmap='gray')\n",
    "        plt.axis('off')\n",
    "\n",
    "        # code\n",
    "        plt.subplot(5, 5, i+2)\n",
    "        plt.imshow(code[i, :, :], cmap='gray')\n",
    "        plt.axis('off')\n",
    "\n",
    "        # output im\n",
    "        plt.subplot(5, 5, i+3)\n",
    "        plt.imshow(construct_img1[i, :,].reshape((28, 28)), cmap='gray')\n",
    "        plt.axis('off')\n",
    "        \n",
    "        # output im\n",
    "        plt.subplot(5, 5, i+4)\n",
    "        plt.imshow(construct_img2[i, :,].reshape((28, 28)), cmap='gray')\n",
    "        plt.axis('off')\n",
    "        \n",
    "        # output im\n",
    "        plt.subplot(5, 5, i+5)\n",
    "        plt.imshow(construct_img3[i, :,].reshape((28, 28)), cmap='gray')\n",
    "        plt.axis('off')\n",
    "\n",
    "    #plt.savefig('vae-pic/vae_rebuilt.png')\n",
    "\n",
    "save_plt(x_test[10:110], logits_x1,  logits_x2,  logits_x3, code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_squares(images, nr_images_per_side):\n",
    "    images_to_plot = np.concatenate(\n",
    "        [np.concatenate([images[j*nr_images_per_side+i].reshape((28,28)) for i in range(0,nr_images_per_side)],\n",
    "                        axis=1)\n",
    "         for j in range(0,nr_images_per_side)],\n",
    "        axis=0)\n",
    "    return images_to_plot\n",
    "\n",
    "def plot_squares(originals, reconstructs, nr_images_per_side):\n",
    "    originals_square = make_squares(originals, nr_images_per_side)\n",
    "    plt.imsave('original.pdf', originals_square, cmap='viridis', format='pdf')\n",
    "    reconstructs_square = make_squares(reconstructs, nr_images_per_side)\n",
    "    plt.imsave('recons.pdf', reconstructs_square, cmap='viridis', format='pdf')\n",
    "    combined = np.concatenate([originals_square, reconstructs_square], axis=1)\n",
    "    plt.imsave('combined.pdf', combined, cmap='viridis', format='pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_squares(x_test[10:110], logits_x1, 8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gumbel",
   "language": "python",
   "name": "gumbel"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}