artistic_alexnet_train_rmsprop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Theano tutorial : 'A Neural Algorithm of Artistic Style'\n",
    "\n",
    "** Author : Sihyeon Seong (sihyun0826@kaist.ac.kr)\n",
    "\n",
    "Section 3 : Generating artistic images with ConvNet\n",
    "\n",
    "- (Construct convolution & pooling layer) + define loss function for measuring similarity \n",
    "- (Define our model & functions)\n",
    "- Define our optimization methods\n",
    "- Showing our image under optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "\n",
    "import os\n",
    "import sys\n",
    "import timeit\n",
    "\n",
    "from IPython.display import display, clear_output\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import theano\n",
    "import theano.tensor as T\n",
    "from theano.tensor.signal import downsample, pool\n",
    "from theano.tensor.nnet import conv2d\n",
    "\n",
    "from pylearn2.expr.normalize import CrossChannelNormalization\n",
    "\n",
    "theano.config.floatX='float32'\n",
    "\n",
    "params_path = 'pretrained_weights/parameters_releasing'\n",
    "\n",
    "\n",
    "class ConvPoolLayer(object):\n",
    "\n",
    "    def __init__(self, input, filter_shape, image_shape, f_params_w, f_params_b, lrn=False, t_style=None, t_content=None, convstride=1, padsize =0, group=1, poolsize = 3, poolstride = 1):\n",
    "\n",
    "        self.input = input\n",
    "\n",
    "        if t_style is not None:\n",
    "            self.t_style = np.asarray(np.load(t_style),dtype=theano.config.floatX)\n",
    "\n",
    "        if t_content is not None:\n",
    "            self.t_content = np.asarray(np.load(t_content),dtype=theano.config.floatX)\n",
    "\n",
    "        if lrn is True:\n",
    "            self.lrn_func = CrossChannelNormalization()\n",
    "\n",
    "        if group == 1:\n",
    "            self.W = theano.shared(np.asarray(np.transpose(np.load(os.path.join(params_path,f_params_w)),(3,0,1,2)),dtype=theano.config.floatX), borrow=True)\n",
    "            self.b = theano.shared(np.asarray(np.load(os.path.join(params_path,f_params_b)),dtype=theano.config.floatX), borrow=True)\n",
    "            conv_out = conv2d(input=self.input,filters=self.W,filter_shape=filter_shape,border_mode = padsize,subsample=(convstride, convstride),filter_flip=True)\n",
    "\n",
    "        elif group == 2:\n",
    "            self.filter_shape = np.asarray(filter_shape)\n",
    "            self.image_shape = np.asarray(image_shape)\n",
    "            self.filter_shape[0] = self.filter_shape[0] / 2\n",
    "            self.filter_shape[1] = self.filter_shape[1] / 2\n",
    "            self.image_shape[1] = self.image_shape[1] / 2\n",
    "            self.W0 = theano.shared(np.asarray(np.transpose(np.load(os.path.join(params_path,f_params_w[0])),(3,0,1,2)),dtype=theano.config.floatX), borrow=True)\n",
    "            self.W1 = theano.shared(np.asarray(np.transpose(np.load(os.path.join(params_path,f_params_w[1])),(3,0,1,2)),dtype=theano.config.floatX), borrow=True)\n",
    "            self.b0 = theano.shared(np.asarray(np.load(os.path.join(params_path,f_params_b[0])),dtype=theano.config.floatX), borrow=True)\n",
    "            self.b1 = theano.shared(np.asarray(np.load(os.path.join(params_path,f_params_b[1])),dtype=theano.config.floatX), borrow=True)\n",
    "            conv_out0 = conv2d(input=self.input[:,:self.image_shape[1],:,:],filters=self.W0,filter_shape=tuple(self.filter_shape),border_mode = padsize,subsample=(convstride, convstride),filter_flip=True) + self.b0.dimshuffle('x', 0, 'x', 'x')\n",
    "            conv_out1 = conv2d(input=self.input[:,self.image_shape[1]:,:,:],filters=self.W1,filter_shape=tuple(self.filter_shape),border_mode = padsize,subsample=(convstride, convstride),filter_flip=True) + self.b1.dimshuffle('x', 0, 'x', 'x')\n",
    "            conv_out = T.concatenate([conv_out0, conv_out1],axis=1)\n",
    "\n",
    "        else:\n",
    "            raise AssertionError()\n",
    "\n",
    "        relu_out = T.maximum(conv_out, 0)\n",
    "        \n",
    "        if poolsize != 1:\n",
    "            self.output = pool.pool_2d(input=relu_out,ds=(poolsize,poolsize),ignore_border=True, st=(poolstride,poolstride),mode='average_exc_pad')\n",
    "        else:\n",
    "            self.output = relu_out\n",
    "\n",
    "        if lrn is True:\n",
    "            self.output = self.lrn_func(self.output)\n",
    "\n",
    "    def style_error(self):\n",
    "        gram_matrix_ori = T.dot(self.t_style.reshape((self.t_style.shape[1],self.t_style.shape[2]*self.t_style.shape[3])),self.t_style.reshape((self.t_style.shape[1],self.t_style.shape[2]*self.t_style.shape[3])).T)\n",
    "        gram_matrix_gen = T.dot(self.output.reshape((self.t_style.shape[1],self.t_style.shape[2]*self.t_style.shape[3])),self.output.reshape((self.t_style.shape[1],self.t_style.shape[2]*self.t_style.shape[3])).T)\n",
    "        return T.sum(T.sum(T.sqr(gram_matrix_gen-gram_matrix_ori)))/(4.0*(self.t_style.shape[1]**2)*((self.t_style.shape[2]*self.t_style.shape[3])**2))\n",
    "\n",
    "    def content_error(self):\n",
    "        return T.sum(T.sum(T.sum(T.sqr(self.output-self.t_content))))/2.0\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#THEANO_FLAGS='floatX=float32,device=cpu,nvcc.fastmath=True' python artistic_alexnet_train_rmsprop.py\n",
    "\n",
    "def evaluate_alexnet(batch_size=1):\n",
    "\n",
    "    rng = np.random.RandomState(23455)\n",
    "\n",
    "    input_img = np.random.normal(0.0, 1.0, size=(1,3,227,227)).astype(np.float32)\n",
    "\n",
    "    lr = T.fscalar('lr')\n",
    "\n",
    "    x = theano.shared(input_img,borrow=True)\n",
    "    print('... building the model')\n",
    "\n",
    "    layer1_input = x.reshape((batch_size, 3, 227, 227))\n",
    "\n",
    "    convpool_layer1 = ConvPoolLayer(input=layer1_input, image_shape=(batch_size, 3, 227, 227), filter_shape=(96, 3, 11, 11), f_params_w='W_0_65.npy', f_params_b='b_0_65.npy', t_style = 'cnn_features/van_gogh_starry_night_1.npy', t_content = 'cnn_features/kaist_n1_1.npy', lrn=True, convstride=4, padsize=0, group=1, poolsize=3, poolstride=2)\n",
    "\n",
    "    convpool_layer2 = ConvPoolLayer(input=convpool_layer1.output,image_shape=(batch_size, 96, 27, 27),filter_shape=(256, 96, 5, 5), f_params_w=['W0_1_65.npy','W1_1_65.npy'], t_style = 'cnn_features/van_gogh_starry_night_2.npy', lrn=True, f_params_b=['b0_1_65.npy','b1_1_65.npy'], convstride=1, padsize=2, group=2, poolsize=3, poolstride=2)\n",
    "\n",
    "    convpool_layer3 = ConvPoolLayer(input=convpool_layer2.output,image_shape=(batch_size, 256, 13, 13),filter_shape=(384, 256, 3, 3), f_params_w='W_2_65.npy', f_params_b='b_2_65.npy', t_style = 'cnn_features/van_gogh_starry_night_3.npy',convstride=1, padsize=1, group=1,poolsize=1, poolstride=0)\n",
    "\n",
    "    convpool_layer4 = ConvPoolLayer(input=convpool_layer3.output,image_shape=(batch_size, 384, 13, 13),filter_shape=(384, 384, 3, 3), f_params_w=['W0_3_65.npy','W1_3_65.npy'], t_style = 'cnn_features/van_gogh_starry_night_4.npy', f_params_b=['b0_3_65.npy','b1_3_65.npy'],convstride=1, padsize=1, group=2,poolsize=1, poolstride=0)\n",
    "\n",
    "    convpool_layer5 = ConvPoolLayer(input=convpool_layer4.output,image_shape=(batch_size, 384, 13, 13),filter_shape=(256, 384, 3, 3), f_params_w=['W0_4_65.npy','W1_4_65.npy'], t_style = 'cnn_features/van_gogh_starry_night_5.npy', f_params_b=['b0_4_65.npy','b1_4_65.npy'], convstride=1, padsize=1, group=2,poolsize=3, poolstride=2)\n",
    "\n",
    "    cost= 0.2*(convpool_layer1.style_error() + convpool_layer2.style_error() +\n",
    "               convpool_layer3.style_error() + convpool_layer4.style_error() +\n",
    "               convpool_layer5.style_error()) + 0.00002*convpool_layer1.content_error()\n",
    "\n",
    "    img_out = theano.function([],x)\n",
    "\n",
    "    print('... train')\n",
    "\n",
    "    params = x\n",
    "\n",
    "    grads = T.grad(cost, params)\n",
    "\n",
    "    #####RMSprop\n",
    "    decay = 0.9\n",
    "    max_scaling=1e5\n",
    "    epsilon = 1. / max_scaling\n",
    "\n",
    "    vels = theano.shared(params.get_value() * 0.) \n",
    "\n",
    "    new_mean_squared_grad = (decay * vels + (1 - decay) * T.sqr(grads))\n",
    "    rms_grad_t = T.sqrt(new_mean_squared_grad)\n",
    "    delta_x_t = - lr * grads / rms_grad_t\n",
    "\n",
    "    updates=[]\n",
    "    updates.append((params,params + delta_x_t))\n",
    "    updates.append((vels,new_mean_squared_grad))\n",
    "\n",
    "\n",
    "    train_model = theano.function([lr],[cost],updates=updates)\n",
    "\n",
    "    n_epochs = 3000\n",
    "\n",
    "    img_mean = np.load('pretrained_weights/img_mean.npy')\n",
    "    img_mean = np.transpose(img_mean,(1,2,0))\n",
    "\n",
    "    tmp = np.transpose(np.squeeze(input_img),(1,2,0))\n",
    "\n",
    "    recon = tmp+img_mean[16:16+227,16:16+227,:]\n",
    "\n",
    "    fig = plt.figure()\n",
    "    fig_handle = plt.imshow(recon.astype(np.uint8))\n",
    "    display(fig)\n",
    "\n",
    "    learning_rate = 1.0\n",
    "    lr_phase = 0\n",
    "    for i in xrange(n_epochs):\n",
    "        img_gen = np.transpose(np.squeeze(img_out()),(1,2,0))\n",
    "        recon = img_gen+img_mean[16:16+227,16:16+227,:]\n",
    "        fig_handle.set_data(recon.astype(np.uint8))\n",
    "        clear_output(wait=True)\n",
    "        display(fig)\n",
    "\n",
    "        print train_model(learning_rate)[0], ' / epochs : ', i\n",
    "\n",
    "    results = img_out()\n",
    "    np.save('results.npy',results)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "if __name__ == '__main__':\n",
    "    evaluate_alexnet()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}