polynomial_neural_network_astroid_generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is an example that synthesizes images based on the polynomial neural networks' paper.\n",
    "\n",
    "The code below (short version of the original code) trains a polynomial generator in a GAN framework.\n",
    "\n",
    "That is, the generator is a polynomial with Fully-connected layers and includes no activation functions.\n",
    "\n",
    "Additional examples can be found in the official repo:\n",
    "\n",
    " https://github.com/grigorisg9gr/polynomial_nets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os.path import join, isdir, isfile\n",
    "import numpy as np\n",
    "from functools import partial\n",
    "# # chainer related imports.\n",
    "try:\n",
    "    import chainer\n",
    "except ImportError as e:\n",
    "    print('This notebook depends on chainer, please install it to proceed.')\n",
    "    raise ImportError(e)\n",
    "import chainer.functions as F\n",
    "import chainer.links as L\n",
    "from chainer import Variable\n",
    "from chainer import training\n",
    "from chainer.training import extensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define the dataset class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_astroid(n_samples, alpha=1, **kwargs):\n",
    "    \"\"\" Create an astroid (see wikipedia's lemma for parametric equations). \"\"\"\n",
    "    t = np.linspace(-alpha * 4, alpha * 4, num=n_samples)\n",
    "    x = alpha / 4 * (3 * np.sin(t) - np.sin(3 * t))\n",
    "    y = alpha / 4 * (3 * np.cos(t) + np.cos(3 * t))\n",
    "    return np.vstack((x, y)).astype(np.float32).T\n",
    "\n",
    "\n",
    "class SyntheticDataset(chainer.dataset.DatasetMixin):\n",
    "    def __init__(self, n_samples, seed=0, **kwargs):\n",
    "        sample_fn = partial(create_astroid, **kwargs)\n",
    "        np.random.seed(seed)\n",
    "        inp_feats = sample_fn(n_samples=n_samples)\n",
    "        perm = np.random.permutation(range(inp_feats.shape[0]))\n",
    "        # # permute and reshape with additional dimensions.\n",
    "        self.base = np.reshape(inp_feats[perm], inp_feats.shape + (1, 1))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.base)\n",
    "\n",
    "    def get_example(self, i):\n",
    "        return self.base[i], 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # # # uncomment below to view the data distribution.\n",
    "# synn1 = SyntheticDataset(300)\n",
    "# pts = np.array([synn1.get_example(i)[0][:, 0, 0] for i in range(len(synn1))])\n",
    "# import matplotlib.pyplot as plt\n",
    "# %matplotlib inline\n",
    "# plt.scatter(pts[:, 0], pts[:, 1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define the generator and discriminator classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FCDiscriminator(chainer.Chain):\n",
    "    def __init__(self, n_input=11, n_out=4, n_hidden=8, n_hidden2=5, sn=False,\n",
    "                 layer_d=None, **kwargs):\n",
    "        # # if layer_d is defined it overrides n_hidden[i] with i >= 3.\n",
    "        # # Note: layer_d should include the value of n_hidden2.\n",
    "        assert (layer_d is None) or n_hidden2 == 10 or layer_d[0] == n_hidden2\n",
    "        assert layer_d is not None\n",
    "        w = chainer.initializers.GlorotUniform()\n",
    "        super(FCDiscriminator, self).__init__()\n",
    "        Linear = SNLinear if sn else L.Linear\n",
    "        self.n_l = n_l = len(layer_d) + 2\n",
    "\n",
    "        with self.init_scope():\n",
    "            self.l1 = Linear(n_input, n_hidden, initialW=w)\n",
    "            self.l2 = Linear(n_hidden, n_hidden2, initialW=w)\n",
    "\n",
    "            # # iterate over all layers (till the last) and save in self.\n",
    "            for l in range(3, n_l):\n",
    "                # # define the input and the output names.\n",
    "                ni, no = layer_d[l - 3], layer_d[l - 3 + 1]\n",
    "                setattr(self, 'l{}'.format(l), Linear(ni, no, initialW=w))\n",
    "\n",
    "            # # save the last layer. \n",
    "            ni = layer_d[n_l - 3]\n",
    "            setattr(self, 'l{}'.format(n_l), Linear(ni, n_out, initialW=w))\n",
    "            # # add the binary classification layer in the end.\n",
    "            self.lin = L.Linear(n_out, 1, initialW=w)\n",
    "        self.activ = F.relu\n",
    "\n",
    "    def __call__(self, x, y=None, **kwargs):\n",
    "        h = x\n",
    "        # # loop over the layers.\n",
    "        for l in range(1, self.n_l):\n",
    "            h = self.activ(getattr(self, 'l{}'.format(l))(h))\n",
    "        # # last layer (no activation).\n",
    "        repres = getattr(self, 'l{}'.format(self.n_l))(h)\n",
    "        return self.lin(repres)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_continuous(dim, batchsize, distribution='normal', xp=np):\n",
    "    if distribution == 'normal':\n",
    "        return xp.random.randn(batchsize, dim).astype(xp.float32)\n",
    "    elif distribution == 'uniform':\n",
    "        return xp.random.uniform(-1, 1, (batchsize, dim)).astype(xp.float32)\n",
    "    else:\n",
    "        raise NotImplementedError\n",
    "\n",
    "\n",
    "class FCProductRecursiveGenerator(chainer.Chain):\n",
    "    def __init__(self, dim_z=2, channels=[3, 10], power_poly=[6, 8], n_out=2, use_bias=True,\n",
    "                 distribution='uniform', use_bn=False, derivfc=0):\n",
    "        \"\"\"\n",
    "        Polynomial generator with fully-connected layers. Specifically, it implements the\n",
    "        product of polynomials with both model 1 of the polygan paper.\n",
    "        :param dim_z: int; the dimensions of the input noise (the prior distribution samples).\n",
    "        :param channels: list; each element (int) includes the depth of the respective FC layer. We consider\n",
    "            that each polynomial has a constant number of layers to avoid obfuscating the code.\n",
    "        :param power_poly: list; each element in the list denotes the power (i.e. approximation N_i)\n",
    "            of each polynomial.\n",
    "        :param n_out: int; the dimensions of the output.\n",
    "        :param use_bias: bool; whether to use bias in the FC layers.\n",
    "        :param distribution: str; the prior distribution.\n",
    "        :param use_bn: bool; whether to use batch normalization.\n",
    "        \"\"\"\n",
    "        w = chainer.initializers.GlorotUniform()\n",
    "        super(FCProductRecursiveGenerator, self).__init__()\n",
    "        Linear = L.Linear\n",
    "        # # Save several attributes from the provided arguments.\n",
    "        assert isinstance(dim_z, int)\n",
    "        self.dim_z = dim_z\n",
    "        self.distribution = distribution\n",
    "        self.use_bn = use_bn\n",
    "        self.channels = channels\n",
    "        # # whether to use bias in the FC layers.\n",
    "        self.use_bias = use_bias\n",
    "        # # the input size to the current polynomial; initialize on dimz.\n",
    "        input_current_poly = dim_z\n",
    "        self.power_poly = power_poly\n",
    "        assert len(self.power_poly) == len(self.channels)\n",
    "        self.derivfc = derivfc\n",
    "        \n",
    "        with self.init_scope():\n",
    "            bn1 = partial(L.BatchNormalization, use_gamma=True, use_beta=False)\n",
    "            # # iterate over all the polynomials (length of channels many).\n",
    "            for id_poly in range(len(self.channels)):\n",
    "                # # set the channels for this polynomial appropriately for each layer.\n",
    "                channels_poly = channels[id_poly]\n",
    "                assert isinstance(channels_poly, int)\n",
    "                # # replace the int with a list (same channels for all). \n",
    "                channels_poly = [channels_poly] * (self.power_poly[id_poly] + 1)\n",
    "                # ensure that the current input channels match the expected.\n",
    "                setattr(self, 'has_rsz{}'.format(id_poly), input_current_poly != channels_poly[0])\n",
    "                if input_current_poly != channels_poly[0]:\n",
    "                    setattr(self, 'resize{}'.format(id_poly), Linear(input_current_poly,\n",
    "                                                                     channels_poly[0], nobias=not use_bias))\n",
    "                # # now build the current polynomial (id_poly).\n",
    "                for l in range(1, self.power_poly[id_poly] + 1):\n",
    "                    c1 = channels_poly[l]\n",
    "                    cin = c1\n",
    "                    setattr(self, 'l{}_{}'.format(id_poly, l), Linear(cin, c1, nobias=not use_bias))\n",
    "                    if use_bn:\n",
    "                        setattr(self, 'bn{}_{}'.format(id_poly, l), bn1(c1))\n",
    "                # # update the channels for the next polynomial.\n",
    "                input_current_poly = int(channels_poly[-1])\n",
    "            # # save the last layer (only for the last). \n",
    "            self.last = Linear(input_current_poly, n_out)\n",
    "\n",
    "    def __call__(self, batchsize, z=None, **kwargs):\n",
    "        if z is None:\n",
    "            z = sample_continuous(self.dim_z, batchsize, distribution=self.distribution, xp=self.xp)\n",
    "        # # input_poly: the input variable to each polynomial; for the \n",
    "        # # first, simply z, i.e. the noise vector.\n",
    "        input_poly = z + 0\n",
    "        # # iterate over all the polynomials (length of channels many).\n",
    "        for id_poly in range(len(self.channels)):\n",
    "            # # ensure that the channels from previous polynomial are of the\n",
    "            # # appropriate size.\n",
    "            if getattr(self, 'has_rsz{}'.format(id_poly)):\n",
    "                input_poly = getattr(self, 'resize{}'.format(id_poly))(input_poly)\n",
    "            h = getattr(self, 'l{}_1'.format(id_poly))(input_poly)\n",
    "            \n",
    "            # # loop over the current polynomial layers and compute the \n",
    "            # # output (for this polynomial). \n",
    "            for layer in range(2, self.power_poly[id_poly] + 1):\n",
    "                # # step 1: perform the hadamard product.\n",
    "                z1 = getattr(self, 'l{}_{}'.format(id_poly, layer))(input_poly)\n",
    "                h = z1 * h + h\n",
    "                # # step 2: normalize representations.\n",
    "                if self.use_bn:\n",
    "                    h = getattr(self, 'bn{}_{}'.format(id_poly, layer))(h)\n",
    "            # # update the input for the next polynomial.\n",
    "            input_poly = h + 0\n",
    "        # # last layer.\n",
    "        h = self.last(h)\n",
    "        if len(h.shape) == 2:\n",
    "            h = F.reshape(h, (h.shape[0], h.shape[1], 1, 1))\n",
    "        return h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define the updater class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hinge Loss\n",
    "def loss_hinge_dis(dis_fake, dis_real):\n",
    "    loss = F.mean(F.relu(1. - dis_real))\n",
    "    loss += F.mean(F.relu(1. + dis_fake))\n",
    "    return loss\n",
    "\n",
    "\n",
    "def loss_hinge_gen(dis_fake):\n",
    "    loss = -F.mean(dis_fake)\n",
    "    return loss\n",
    "\n",
    "\n",
    "class Updater(chainer.training.StandardUpdater):\n",
    "    def __init__(self, gen, dis, opt_gen, opt_dis, iterator, *args, n_dis=1, \n",
    "                 n_gen_samples=512, **kwargs):\n",
    "        self.loss_dis = loss_hinge_dis\n",
    "        self.loss_gen = loss_hinge_gen\n",
    "        self.n_gen_samples = n_gen_samples\n",
    "        self.gen, self.dis = gen, dis\n",
    "        self.opt_gen, self.opt_dis = opt_gen, opt_dis\n",
    "        self.iterator = iterator\n",
    "        self.n_dis = n_dis\n",
    "        kwargs1 = {'iterator': iterator, 'optimizer': {'opt_gen': opt_gen, 'opt_dis': opt_dis}}\n",
    "        super(Updater, self).__init__(*args, **kwargs1)\n",
    "\n",
    "    def _generate_samples(self, n_gen_samples=None):\n",
    "        if n_gen_samples is None:\n",
    "            n_gen_samples = self.n_gen_samples\n",
    "        x_fake = self.gen(n_gen_samples)\n",
    "        return x_fake\n",
    "\n",
    "    def get_batch(self, xp):\n",
    "        batch = self.iterator.next()\n",
    "        x = []\n",
    "        for j in range(len(batch)):\n",
    "            x.append(np.asarray(batch[j][0]).astype('f'))\n",
    "        x_real = Variable(xp.asarray(x))\n",
    "        return x_real\n",
    "\n",
    "    def update_core(self):\n",
    "        gen = self.gen\n",
    "        dis = self.dis\n",
    "        xp = gen.xp\n",
    "        for i in range(self.n_dis):\n",
    "            x_real = self.get_batch(xp)\n",
    "            batchsize = len(x_real)\n",
    "            dis_real = dis(x_real)\n",
    "            x_fake = self._generate_samples(n_gen_samples=batchsize)\n",
    "            dis_fake = dis(x_fake)\n",
    "            x_fake.unchain_backward()\n",
    "\n",
    "            fake_arr, real_arr = dis_fake.array, dis_real.array\n",
    "            chainer.reporter.report({'dis_fake': fake_arr.mean()})\n",
    "            chainer.reporter.report({'dis_real': real_arr.mean()})\n",
    "\n",
    "            loss_dis = self.loss_dis(dis_fake=dis_fake, dis_real=dis_real)\n",
    "            dis.cleargrads()\n",
    "            loss_dis.backward()\n",
    "            self.opt_dis.update()\n",
    "            loss_dis.unchain_backward()\n",
    "            chainer.reporter.report({'loss_dis': loss_dis.array})\n",
    "            del loss_dis\n",
    "\n",
    "            if i == 0:\n",
    "                x_fake = self._generate_samples()\n",
    "                dis_fake = dis(x_fake)\n",
    "                loss_gen = self.loss_gen(dis_fake=dis_fake)\n",
    "                assert not xp.isnan(loss_gen.data)\n",
    "                gen.cleargrads()\n",
    "                loss_gen.backward()\n",
    "                self.opt_gen.update()\n",
    "                loss_gen.unchain_backward()\n",
    "                chainer.reporter.report({'loss_gen': loss_gen.array})\n",
    "                del loss_gen\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare the models for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = 128\n",
    "n_iters = 30000\n",
    "cuda = chainer.cuda.available\n",
    "\n",
    "dataset = SyntheticDataset(n_samples=200000)\n",
    "iterator = chainer.iterators.SerialIterator(dataset, batch)\n",
    "\n",
    "# # define the instances of the generator and the discriminator.\n",
    "dis = FCDiscriminator(n_input=2, n_out=2, n_hidden=20, n_hidden2=20, layer_d=[20, 20, 20, 20, 20])\n",
    "gen = FCProductRecursiveGenerator(channels=[3, 10], power_poly=[6, 8])\n",
    "\n",
    "if cuda:\n",
    "    gen.to_gpu()\n",
    "    dis.to_gpu()\n",
    "\n",
    "# # define the optimizers for the generator and the discriminator.\n",
    "opt_gen = chainer.optimizers.Adam(alpha=0.00015, beta1=0., beta2=0.9)\n",
    "opt_gen.setup(gen)\n",
    "opt_dis = chainer.optimizers.Adam(alpha=0.0001, beta1=0., beta2=0.9)\n",
    "opt_dis.setup(dis)\n",
    "\n",
    "# # initialize the trainer and the updater.\n",
    "updater = Updater(gen, dis, opt_gen, opt_dis, iterator, n_dis=3)\n",
    "trainer = training.Trainer(updater, (n_iters, 'iteration'),) # out=out\n",
    "# # add extensions to the training (auxiliary for the user).\n",
    "trainer.extend(extensions.LogReport(trigger=(200, 'iteration')))\n",
    "report_keys = ['loss_dis', 'loss_gen', 'synth_qual']\n",
    "trainer.extend(extensions.PrintReport(report_keys), trigger=(200, 'iteration'))\n",
    "trainer.extend(extensions.ProgressBar(update_interval=200))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m1 = 'Generator params: {}. Discriminator params: {}.'\n",
    "print(m1.format(gen.count_params(), dis.count_params()))\n",
    "# Run the training\n",
    "print('start training')\n",
    "trainer.run()\n",
    "# # export the last model.\n",
    "extensions.snapshot_object(gen, '{}_best.npz'.format(gen.__class__.__name__))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Synthesize points from the trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aa = gen(10000).array\n",
    "plt.scatter(aa[:, 0, 0, 0], aa[:, 1, 0, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}