pix2pix-lasagne.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['THEANO_FLAGS']='floatX=float32,device=cuda,optimizer=fast_run,dnn.library_path=/usr/lib'\n",
    "channel_first = True\n",
    "channel_axis=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import theano\n",
    "import theano.tensor as T\n",
    "import lasagne\n",
    "from lasagne.layers import DropoutLayer, ReshapeLayer, InputLayer \n",
    "floatX = theano.config.floatX\n",
    "from lasagne.layers import Conv2DLayer, TransposedConv2DLayer, ConcatLayer, NonlinearityLayer\n",
    "from lasagne.layers import batch_norm\n",
    "from lasagne.nonlinearities import LeakyRectify, sigmoid, rectify, tanh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "conv_init = lasagne.init.Normal(0.02, 0)\n",
    "gamma_init = lasagne.init.Normal(0.02, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def BASIC_D(nc_in, nc_out, ndf, max_layers=3, use_sigmoid=True):\n",
    "    l = -1\n",
    "    def conv2d(x, nf, stride=2, nonlinearity=LeakyRectify(0.2)):\n",
    "        nonlocal l\n",
    "        l+=1\n",
    "        return Conv2DLayer(x, num_filters=nf, filter_size=4, stride=stride, \n",
    "                           pad=1, W=conv_init, flip_filters=False,\n",
    "                            nonlinearity=nonlinearity,\n",
    "                           name=\"conv2d_{}\".format(l)\n",
    "                          )\n",
    "    input_a = InputLayer(shape=(None, nc_in, None, None), name=\"inputA\")\n",
    "    input_b = InputLayer(shape=(None, nc_out, None, None), name=\"inputB\")\n",
    "    _ = ConcatLayer([input_a, input_b], name='concat')\n",
    "    _ = conv2d(_, ndf)\n",
    "    for layer in range(1, max_layers):        \n",
    "        out_feat = ndf * min(2**layer, 8)\n",
    "        _ = conv2d(_, out_feat)\n",
    "        _ = batch_norm(_, epsilon=1e-5, gamma=gamma_init)\n",
    "    out_feat = ndf*min(2**max_layers, 8)\n",
    "    _ = conv2d(_, out_feat, stride=1)\n",
    "    _ = batch_norm(_, epsilon=1e-5, gamma=gamma_init)         \n",
    "    _ = conv2d(_, 1, stride=1, nonlinearity=sigmoid if use_sigmoid else None)\n",
    "    return _\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# from https://gist.github.com/ajbrock/a3858c26282d9731191901b397b3ce9f\n",
    "def reflect_pad(x, width, batch_ndim=1):\n",
    "    \"\"\"\n",
    "    Pad a tensor with a constant value.\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : tensor\n",
    "    width : int, iterable of int, or iterable of tuple\n",
    "        Padding width. If an int, pads each axis symmetrically with the same\n",
    "        amount in the beginning and end. If an iterable of int, defines the\n",
    "        symmetric padding width separately for each axis. If an iterable of\n",
    "        tuples of two ints, defines a seperate padding width for each beginning\n",
    "        and end of each axis.\n",
    "    batch_ndim : integer\n",
    "        Dimensions before the value will not be padded.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Idea for how to make this happen: Flip the tensor horizontally to grab horizontal values, then vertically to grab vertical values\n",
    "    # alternatively, just slice correctly\n",
    "    input_shape = x.shape\n",
    "    input_ndim = x.ndim\n",
    "\n",
    "    output_shape = list(input_shape)\n",
    "    indices = [slice(None) for _ in output_shape]\n",
    "\n",
    "    if isinstance(width, int):\n",
    "        widths = [width] * (input_ndim - batch_ndim)\n",
    "    else:\n",
    "        widths = width\n",
    "\n",
    "    for k, w in enumerate(widths):\n",
    "        try:\n",
    "            l, r = w\n",
    "        except TypeError:\n",
    "            l = r = w\n",
    "        output_shape[k + batch_ndim] += l + r\n",
    "        indices[k + batch_ndim] = slice(l, l + input_shape[k + batch_ndim])\n",
    "\n",
    "    # Create output array\n",
    "    out = T.zeros(output_shape)\n",
    "    \n",
    "    # Vertical Reflections\n",
    "    out=T.set_subtensor(out[:,:,:width,width:-width], x[:,:,width:0:-1,:])# out[:,:,:width,width:-width] = x[:,:,width:0:-1,:]\n",
    "    out=T.set_subtensor(out[:,:,-width:,width:-width], x[:,:,-2:-(2+width):-1,:])#out[:,:,-width:,width:-width] = x[:,:,-2:-(2+width):-1,:]\n",
    "    \n",
    "    # Place X in out\n",
    "    # out = T.set_subtensor(out[tuple(indices)], x) # or, alternative, out[width:-width,width:-width] = x\n",
    "    out=T.set_subtensor(out[:,:,width:-width,width:-width],x)#out[:,:,width:-width,width:-width] = x\n",
    "   \n",
    "   #Horizontal reflections\n",
    "    out=T.set_subtensor(out[:,:,:,:width],out[:,:,:,(2*width):width:-1])#out[:,:,:,:width] = out[:,:,:,(2*width):width:-1]\n",
    "    out=T.set_subtensor(out[:,:,:,-width:],out[:,:,:,-(width+2):-(2*width+2):-1])#out[:,:,:,-width:] = out[:,:,:,-(width+2):-(2*width+2):-1]\n",
    "    \n",
    "    \n",
    "    return out\n",
    "    \n",
    "class ReflectLayer(lasagne.layers.Layer):\n",
    "\n",
    "    def __init__(self, incoming, width, batch_ndim=2, **kwargs):\n",
    "        super(ReflectLayer, self).__init__(incoming, **kwargs)\n",
    "        self.width = width\n",
    "        self.batch_ndim = batch_ndim\n",
    "\n",
    "    def get_output_shape_for(self, input_shape):\n",
    "        output_shape = list(input_shape)\n",
    "\n",
    "        if isinstance(self.width, int):\n",
    "            widths = [self.width] * (len(input_shape) - self.batch_ndim)\n",
    "        else:\n",
    "            widths = self.width\n",
    "\n",
    "        for k, w in enumerate(widths):\n",
    "            if output_shape[k + self.batch_ndim] is None:\n",
    "                continue\n",
    "            else:\n",
    "                try:\n",
    "                    l, r = w\n",
    "                except TypeError:\n",
    "                    l = r = w\n",
    "                output_shape[k + self.batch_ndim] += l + r\n",
    "        return tuple(output_shape)\n",
    "\n",
    "    def get_output_for(self, input, **kwargs):\n",
    "        return reflect_pad(input, self.width, self.batch_ndim) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def UNET_G(isize, nc_in=3, nc_out=3, ngf=64, fixed_input_size=True):    \n",
    "    max_nf = 8*ngf    \n",
    "    def block(x, s, nf_in, use_batchnorm=True, nf_out=None, nf_next=None):\n",
    "        # print(\"block\",x,s,nf_in, use_batchnorm, nf_out, nf_next)\n",
    "        assert s>=2 and s%2==0\n",
    "        if nf_next is None:\n",
    "            nf_next = min(nf_in*2, max_nf)\n",
    "        if nf_out is None:\n",
    "            nf_out = nf_in\n",
    "            \n",
    "        x = Conv2DLayer(x, num_filters=nf_next, filter_size=4, stride=2, pad=1, W=conv_init, flip_filters=False,                \n",
    "                nonlinearity=None, name='conv2d_{}'.format(s))\n",
    "        if s>2:\n",
    "            if use_batchnorm:\n",
    "                x = batch_norm(x, epsilon=1e-5, gamma=gamma_init)\n",
    "            x2 = NonlinearityLayer(x, nonlinearity=LeakyRectify(0.2), name=\"leakyRelu_{}\".format(s))\n",
    "            x2 = block(x2, s//2, nf_next)\n",
    "            x = ConcatLayer([x, x2], name=\"concat_{}\".format(s))            \n",
    "        x = NonlinearityLayer(x, nonlinearity=rectify, name=\"Relu_{}\".format(s))\n",
    "        x = TransposedConv2DLayer(x, num_filters=nf_out, filter_size=4, stride=2, crop=1, W=conv_init, \n",
    "                                  flip_filters=True, nonlinearity=None, name=\"convt_{}\".format(s))\n",
    "        if use_batchnorm:\n",
    "            x = batch_norm(x, epsilon=1e-5, gamma=gamma_init)\n",
    "        if s <= 8:\n",
    "            x = DropoutLayer(x, 0.5, name=\"dropout_{}\".format(s))\n",
    "        return x\n",
    "    \n",
    "    s = isize if fixed_input_size else None\n",
    "    _ = InputLayer(shape=(None, nc_in, s, s), name='input')\n",
    "    _ = block(_, isize, nc_in, False, nf_out=nc_out, nf_next=ngf)\n",
    "    _ = NonlinearityLayer(_, nonlinearity=tanh, name='tanh')\n",
    "    return _"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from lasagne.layers import ElemwiseSumLayer, SliceLayer\n",
    "def reflect_padding_conv(_, num_filters, filter_size=3, stride=1, nonlinearity=rectify, use_batchnorm=True, **k):\n",
    "    assert filter_size%2==1\n",
    "    pad_size = filter_size>>1\n",
    "    _ = ReflectLayer(_, width=pad_size)\n",
    "    _ = Conv2DLayer(_, num_filters=num_filters, filter_size=filter_size, stride=stride, \n",
    "                           pad=0, W=conv_init, flip_filters=False, nonlinearity=nonlinearity, **k)\n",
    "    if use_batchnorm:\n",
    "        _ = batch_norm(_, epsilon=1e-5, gamma=gamma_init)\n",
    "    return _\n",
    "def res_block(_, num_filters, name):\n",
    "    x = _\n",
    "    _ = reflect_padding_conv(_, num_filters, name=name+\"_conv1\")\n",
    "    _ = reflect_padding_conv(_, num_filters, nonlinearity=None, name=name+\"_conv2\")\n",
    "    return ElemwiseSumLayer([x, _], name=name+\"_add\")\n",
    "\n",
    "def RESNET_G(isize, nc_in=3, nc_out=3, ngf=64, fixed_input_size=True):\n",
    "    s = isize if fixed_input_size else None    \n",
    "    _ = InputLayer(shape=(None, nc_in, s, s), name='input')    \n",
    "    _ = reflect_padding_conv(_, ngf, 7, name=\"first\")    \n",
    "    for m in [2,4]:\n",
    "        _ = Conv2DLayer(_, num_filters=ngf*m, filter_size=4, stride=2, \n",
    "                           pad=1, W=conv_init, flip_filters=False, \n",
    "                           nonlinearity=rectify, name='conv_{}'.format(ngf*m))\n",
    "        _ = batch_norm(_, epsilon=1e-5, gamma=gamma_init)\n",
    "    for i in range(6):\n",
    "        _ = res_block(_, ngf*4, \"res_block{}\".format(i))\n",
    "    for m in [2,1]:\n",
    "        _ = TransposedConv2DLayer(_, num_filters=ngf*m, filter_size=3, stride=2, \n",
    "                            crop=0, W=conv_init, flip_filters=True,\n",
    "                            nonlinearity=rectify, name=\"convt_{}\".format(ngf*m))\n",
    "        _ = batch_norm(_, epsilon=1e-5, gamma=gamma_init)\n",
    "        _ = SliceLayer(_, slice(0, -1),2)\n",
    "        _ = SliceLayer(_, slice(0, -1),3)    \n",
    "    _ = reflect_padding_conv(_, nc_out, 7, nonlinearity=tanh, use_batchnorm=False, name=\"output\")\n",
    "    return _"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "nc_in = 3\n",
    "nc_out = 3\n",
    "ngf = 64\n",
    "ndf = 64\n",
    "use_lsgan = True\n",
    "λ = 10 if use_lsgan else 100\n",
    "\n",
    "loadSize = 286\n",
    "imageSize = 256\n",
    "batchSize = 1\n",
    "lrD = 2e-4\n",
    "lrG = 2e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "netD = BASIC_D(nc_in, nc_out, ndf, use_sigmoid = not use_lsgan)\n",
    "for l in lasagne.layers.get_all_layers(netD):\n",
    "    print(l.name,  l.output_shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#netG = UNET_G(imageSize, nc_in, nc_out, ngf)\n",
    "netG = RESNET_G(imageSize, nc_in, nc_out, ngf)\n",
    "\n",
    "for l in lasagne.layers.get_all_layers(netG):\n",
    "    print(l.name,  l.output_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from lasagne.layers import get_output, get_all_layers,get_all_params\n",
    "no_bn_avg = dict( batch_norm_update_averages=False,\n",
    "                       batch_norm_use_averages=False)\n",
    "real_A = get_all_layers(netG)[0].input_var\n",
    "fake_B = get_output(netG, **no_bn_avg)\n",
    "netG_generate =  theano.function([real_A], fake_B)\n",
    "\n",
    "netD_l1, netD_l2 = get_all_layers(netD)[:2]\n",
    "real_B = netD_l2.input_var\n",
    "\n",
    "output_D_real = get_output(netD, inputs={netD_l1: real_A, netD_l2: real_B}, **no_bn_avg)\n",
    "output_D_fake = get_output(netD, inputs={netD_l1: real_A, netD_l2: fake_B}, **no_bn_avg)\n",
    "if use_lsgan:\n",
    "    loss_fn = lambda output, target : T.mean( (output-target)**2 )\n",
    "else:\n",
    "    loss_fn = lambda output, target : -T.mean(T.log(output+1e-12)*target+T.log(1-output+1e-12)*(1-target))\n",
    "\n",
    "loss_D_real = loss_fn(output_D_real, T.ones_like(output_D_real))\n",
    "loss_D_fake = loss_fn(output_D_fake, T.zeros_like(output_D_fake))\n",
    "loss_D = loss_D_real + loss_D_fake\n",
    "\n",
    "loss_G_fake = loss_fn(output_D_fake, T.ones_like(output_D_fake))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "loss_L1 = T.mean(abs(fake_B-real_B))\n",
    "loss_G =  loss_G_fake + λ * loss_L1\n",
    "\n",
    "params_netD = get_all_params(netD, trainable=True) \n",
    "params_netG = get_all_params(netG, trainable=True)\n",
    "\n",
    "\n",
    "optimize_G = lasagne.updates.adam(loss_G, params_netG, learning_rate=lrG, beta1=0.5)\n",
    "optimize_D = lasagne.updates.adam(loss_D, params_netD, learning_rate=lrD, beta1=0.5)\n",
    "netG_train =  theano.function([real_A, real_B], [loss_G_fake, loss_L1], updates=optimize_G)\n",
    "netD_train = theano.function([real_A, real_B], [loss_D/2], updates=optimize_D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import numpy as np\n",
    "import glob\n",
    "from random import randint, shuffle\n",
    "\n",
    "def load_data(file_pattern):\n",
    "    return glob.glob(file_pattern)\n",
    "def read_image(fn, direction=0):\n",
    "    im = Image.open(fn)\n",
    "    im = im.resize( (loadSize*2, loadSize), Image.BILINEAR )\n",
    "    arr = np.array(im)/255*2-1\n",
    "    w1,w2 = (loadSize-imageSize)//2,(loadSize+imageSize)//2\n",
    "    h1,h2 = w1,w2\n",
    "    imgA = arr[h1:h2, loadSize+w1:loadSize+w2, :]\n",
    "    imgB = arr[h1:h2, w1:w2, :]\n",
    "    if randint(0,1):\n",
    "        imgA=imgA[:,::-1]\n",
    "        imgB=imgB[:,::-1]\n",
    "    if channel_first:\n",
    "        imgA = np.moveaxis(imgA, 2, 0)\n",
    "        imgB = np.moveaxis(imgB, 2, 0)\n",
    "    if direction==0:\n",
    "        return imgA, imgB\n",
    "    else:\n",
    "        return imgB,imgA\n",
    "\n",
    "data = \"edges2shoes\"\n",
    "#data = \"facades\"\n",
    "direction = 1\n",
    "trainAB = load_data('pix2pix/{}/train/*.jpg'.format(data))\n",
    "valAB = load_data('pix2pix/{}/val/*.jpg'.format(data))\n",
    "assert len(trainAB) and len(valAB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def minibatch(dataAB, batchsize, direction=0):\n",
    "    length = len(dataAB)\n",
    "    epoch = i = 0\n",
    "    tmpsize = None    \n",
    "    while True:\n",
    "        size = tmpsize if tmpsize else batchsize\n",
    "        if i+size > length:\n",
    "            shuffle(dataAB)\n",
    "            i = 0\n",
    "            epoch+=1        \n",
    "        dataA = []\n",
    "        dataB = []\n",
    "        for j in range(i,i+size):\n",
    "            imgA,imgB = read_image(dataAB[j], direction)\n",
    "            dataA.append(imgA)\n",
    "            dataB.append(imgB)\n",
    "        dataA = np.float32(dataA)\n",
    "        dataB = np.float32(dataB)\n",
    "        i+=size\n",
    "        tmpsize = yield epoch, dataA, dataB        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from IPython.display import display\n",
    "def showX(X, rows=1):\n",
    "    assert X.shape[0]%rows == 0\n",
    "    int_X = ( (X+1)/2*255).clip(0,255).astype('uint8')\n",
    "    if channel_first:\n",
    "        int_X = np.moveaxis(int_X.reshape(-1,3,imageSize,imageSize), 1, 3)\n",
    "    else:\n",
    "        int_X = int_X.reshape(-1,imageSize,imageSize, 3)\n",
    "    int_X = int_X.reshape(rows, -1, imageSize, imageSize,3).swapaxes(1,2).reshape(rows*imageSize,-1, 3)\n",
    "    display(Image.fromarray(int_X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "train_batch = minibatch(trainAB, 12, direction=direction)\n",
    "_, trainA, trainB = next(train_batch)\n",
    "showX(trainA, 2)\n",
    "showX(trainB, 2)\n",
    "del train_batch, trainA, trainB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def netG_gen(A):\n",
    "    return np.concatenate([netG_generate(A[i:i+1]) for i in range(A.shape[0])], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from IPython.display import clear_output\n",
    "t0 = time.time()\n",
    "niter = 150\n",
    "gen_iterations = 0\n",
    "errL1 = epoch = errG = 0\n",
    "errL1_sum = errG_sum = errD_sum = 0\n",
    "\n",
    "display_iters = 500\n",
    "val_batch = minibatch(valAB, 6, direction)\n",
    "train_batch = minibatch(trainAB, batchSize, direction)\n",
    "\n",
    "while epoch < niter: \n",
    "    epoch, trainA, trainB = next(train_batch)        \n",
    "    errD,  = netD_train(trainA, trainB)\n",
    "    errD_sum +=errD\n",
    "\n",
    "    # epoch, trainA, trainB = next(train_batch)\n",
    "    errG, errL1 = netG_train(trainA, trainB)\n",
    "    errG_sum += errG\n",
    "    errL1_sum += errL1\n",
    "    gen_iterations+=1\n",
    "    if gen_iterations%display_iters==0:\n",
    "        if gen_iterations%(5*display_iters)==0:\n",
    "            clear_output()\n",
    "        print('[%d/%d][%d] Loss_D: %f Loss_G: %f loss_L1: %f'\n",
    "        % (epoch, niter, gen_iterations, errD_sum/display_iters, errG_sum/display_iters, errL1_sum/display_iters), time.time()-t0)\n",
    "        _, valA, valB = train_batch.send(6) \n",
    "        fakeB = netG_gen(valA)\n",
    "        showX(np.concatenate([valA, valB, fakeB], axis=0), 3)\n",
    "        errL1_sum = errG_sum = errD_sum = 0\n",
    "        _, valA, valB = next(val_batch)\n",
    "        fakeB = netG_gen(valA)\n",
    "        showX(np.concatenate([valA, valB, fakeB], axis=0), 3)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}