GRAMFlow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RM-MMDnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from ops import *\n",
    "from utils import *\n",
    "import os\n",
    "import time\n",
    "from glob import glob\n",
    "from scipy.misc import imsave as ims\n",
    "from random import randint\n",
    "from data_providers import *\n",
    "slim = tf.contrib.slim\n",
    "import scipy as sp\n",
    "\n",
    "cifar = True\n",
    "%pylab inline\n",
    "\n",
    "cifar = tf.keras.datasets.cifar10.load_data()[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def transform_local(image, npx=64, is_crop=False):\n",
    "    if is_crop:\n",
    "        cropped_image = center_crop(image, npx)\n",
    "    else:\n",
    "        cropped_image = image\n",
    "    return np.array(cropped_image)/127.5 - 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = {\n",
    "    'batch_size':64,\n",
    "    'image_dim':32*32*3,\n",
    "    'c':3,\n",
    "    'h':32,\n",
    "    'w':32\n",
    "}\n",
    "\n",
    "def makeScaleMatrix(num_gen, num_orig):\n",
    "\n",
    "    # first 'N' entries have '1/N', next 'M' entries have '-1/M'\n",
    "    s1 =  tf.constant(1.0 / num_gen, shape = [num_gen, 1])\n",
    "    s2 = -tf.constant(1.0 / num_orig, shape = [num_orig, 1])\n",
    "\n",
    "    return tf.concat([s1, s2],0)\n",
    "\n",
    "def computeRatio(x, gen_x, sigma = [1,2,4,8,16]):\n",
    "    \n",
    "    x = slim.flatten(x)\n",
    "    gen_x = slim.flatten(gen_x)\n",
    "\n",
    "    # concatenation of the generated images and images from the dataset\n",
    "    # first 'N' rows are the generated ones, next 'M' are from the data\n",
    "    X = tf.concat([gen_x, x],0)\n",
    "\n",
    "    # dot product between all combinations of rows in 'X'\n",
    "    XX = tf.matmul(X, tf.transpose(X))\n",
    "\n",
    "    # dot product of rows with themselves\n",
    "    X2 = tf.reduce_sum(X * X, 1, keep_dims = True)\n",
    "\n",
    "    # exponent entries of the RBF kernel (without the sigma) for each\n",
    "    # combination of the rows in 'X'\n",
    "    # -0.5 * (x^Tx - 2*x^Ty + y^Ty)\n",
    "    exponent = XX - 0.5 * X2 - 0.5 * tf.transpose(X2)\n",
    "\n",
    "    # scaling constants for each of the rows in 'X'\n",
    "    s = makeScaleMatrix(params['batch_size'], params['batch_size'])\n",
    "\n",
    "    # scaling factors of each of the kernel values, corresponding to the\n",
    "    # exponent values\n",
    "    S = tf.matmul(s, tf.transpose(s))\n",
    "\n",
    "    loss1 = 0\n",
    "    loss2 = 0\n",
    "    mmd = 0\n",
    "\n",
    "    # for each bandwidth parameter, compute the MMD value and add them all\n",
    "    n = params['batch_size']\n",
    "    n_sq = float(n*n)\n",
    "    for i in range(len(sigma)):\n",
    "\n",
    "        # kernel values for each combination of the rows in 'X' \n",
    "        kernel_val = tf.exp(1.0 / sigma[i] * exponent)\n",
    "        \n",
    "        k_de_de = kernel_val[:n,:n]\n",
    "        k_de_nu = kernel_val[:n,n:]\n",
    "        k_nu_de = kernel_val[n:,:n]\n",
    "        k_nu_nu = kernel_val[n:,n:]\n",
    "        \n",
    "        loss1 += tf.matmul(tf.matmul(tf.matrix_inverse(k_de_de),k_de_nu),tf.ones(shape=(params['batch_size'],1)))\n",
    "        mmd += tf.reduce_sum(S * kernel_val)\n",
    "\n",
    "\n",
    "    return loss1/float(len(sigma)), tf.sqrt(mmd)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def discriminator(image, reuse=False, i=0):\n",
    "    \n",
    "    with tf.variable_scope(\"disc\") as scope:\n",
    "        if reuse:\n",
    "                scope.reuse_variables()\n",
    "        h=32\n",
    "        w=32\n",
    "        h0 = image\n",
    "        h0 = lrelu(conv2d(h0, 3, df_dim, name='d_h0_conv')) #16x16x32\n",
    "        h1 = lrelu(tf.contrib.layers.batch_norm(conv2d(h0, df_dim, df_dim*2, name='d_h1_conv'+str(i)))) #8x8x64\n",
    "        h2 = lrelu(tf.contrib.layers.batch_norm(conv2d(h1, df_dim*2, df_dim*2*2, name='d_h2_conv'+str(i)))) #4x4x128\n",
    "        h3 = tf.reshape(h2, [batchsize, -1])\n",
    "        h4 = dense(h3, 4*4*df_dim*4, 128, scope='d_h4_lin'+str(i)) #2048\n",
    "        return h4\n",
    "    \n",
    "\n",
    "        \n",
    "    \n",
    "\n",
    "        \n",
    "def generator(z):\n",
    "    with tf.variable_scope(\"gen\") as scope:\n",
    "        gf_dim=32\n",
    "        h=32\n",
    "        w=32\n",
    "        z2 = dense(z, z_dim, 4*4*gf_dim*4, scope='g_h0_lin')\n",
    "        h0 = tf.nn.relu(tf.contrib.layers.batch_norm(tf.reshape(z2, [-1, 4, 4, gf_dim*4]))) # 4x4x128\n",
    "        h1 = tf.nn.relu(tf.contrib.layers.batch_norm(conv_transpose(h0, [batchsize, 8, 8, gf_dim*2], \"g_h1\"))) #8x8x64\n",
    "        h2 = tf.nn.relu(tf.contrib.layers.batch_norm(conv_transpose(h1, [batchsize, 16, 16, gf_dim*1], \"g_h2\"))) #16x16x32\n",
    "        h3 = tf.nn.tanh(conv_transpose(h2, [batchsize, 32, 32, 3], \"g_h4\"))\n",
    "        return h3  \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "sdata=[]\n",
    "batchsize = params['batch_size']\n",
    "iscrop = False\n",
    "imageshape = [32, 32, 3]\n",
    "z_dim = 128\n",
    "gf_dim = 32\n",
    "df_dim = 32\n",
    "c_dim = 3\n",
    "learningrate_gen = 1e-4\n",
    "learningrate = 1e-4\n",
    "beta1 = 0.5\n",
    "NUM_PROJ=1\n",
    "NUM_MMD = 10\n",
    "\n",
    "images = tf.placeholder(tf.float32, [batchsize] + imageshape, name=\"real_images\")\n",
    "\n",
    "zin = tf.placeholder(tf.float32, [None, z_dim], name=\"z\")\n",
    "G = generator(zin)\n",
    "\n",
    "# GAN\n",
    "dloss=0.0\n",
    "gloss=0.0\n",
    "closs=0.0\n",
    "\n",
    "DL = discriminator(images,i=1)\n",
    "GL = discriminator(G,reuse=True,i=1)\n",
    "\n",
    "\n",
    "\n",
    "r_nu,mmd = computeRatio(GL, DL) #correct\n",
    "r_de,_ = computeRatio(DL, GL)\n",
    "\n",
    "gloss = tf.reduce_mean(mmd)\n",
    "\n",
    "dloss = -tf.reduce_mean((r_nu))\n",
    "                      \n",
    "t_vars = tf.trainable_variables()\n",
    "\n",
    "g_vars = [var for var in t_vars if 'gen' in var.name]\n",
    "d_vars = [var for var in t_vars if 'disc' in var.name]\n",
    "\n",
    "print(np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()]))\n",
    "print(np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables() if 'disc' in v.name]))\n",
    "print(np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables() if 'gen' in v.name]))\n",
    "\n",
    "g_optim = tf.train.AdamOptimizer(learningrate_gen, beta1=beta1).minimize(gloss, var_list=g_vars)\n",
    "d_optim = tf.train.AdamOptimizer(learningrate,     beta1=beta1).minimize(dloss, var_list=d_vars)\n",
    "\n",
    "\n",
    "start_time = time.time()\n",
    "display_z = np.random.uniform(-1, 1, [batchsize, z_dim]).astype(np.float32)\n",
    "\n",
    "\n",
    "sess = tf.InteractiveSession()\n",
    "sess.run(tf.global_variables_initializer())\n",
    "tf.initialize_all_variables().run()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shuffle(a):\n",
    "    p = numpy.random.permutation(len(a))\n",
    "    return a[p]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# '------Training GAN--------------------'\n",
    "from tqdm import tqdm_notebook\n",
    "\n",
    "for epoch in xrange(200):\n",
    "    cifar = shuffle(cifar)\n",
    "    \n",
    "    batch_idxs = len(cifar) // params['batch_size']\n",
    "    \n",
    "    for idx in tqdm_notebook(range(batch_idxs)):\n",
    "        batch_images= transform_local(cifar[idx*params['batch_size']:(idx+1)*params['batch_size']])\n",
    "        batch_z = np.random.uniform(-1, 1, [batchsize, z_dim]).astype(np.float32) \n",
    "        \n",
    "        for i in range(1):\n",
    "            sess.run([d_optim],\n",
    "                 feed_dict={ images: batch_images, zin: batch_z })\n",
    "        for i in range(1):\n",
    "            sess.run([g_optim],\n",
    "                     feed_dict={ images: batch_images, zin: batch_z })\n",
    "        \n",
    "#     '---------Printing intermediate results-------------'      \n",
    "    if epoch % 1 == 0:\n",
    "                   \n",
    "        print(\"Epoch: [%2d] [%4d/%4d] time: %4.4f, \" % (epoch, idx, batch_idxs, time.time() - start_time,))\n",
    "        print(dloss.eval({zin: batch_z,images: batch_images}))\n",
    "        \n",
    "        _,sdata = sess.run([images,G],feed_dict={ images: batch_images, zin: batch_z })\n",
    "        sdata = sdata[:64]\n",
    "        sdata = np.expand_dims(sdata,0)\n",
    "        img = merge(sdata[0],[8,8])\n",
    "        plt.imshow(img)\n",
    "        plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}