Inference_firstTry.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-01-19T20:52:00.593277Z",
     "start_time": "2018-01-19T20:52:00.104620Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd import Variable\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import torch.optim as optim\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-01-19T22:44:45.402633Z",
     "start_time": "2018-01-19T22:44:45.229352Z"
    },
    "run_control": {
     "marked": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total images <= max_seq_len is 200\n",
      "total images <= max_seq_len is 100\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import utils\n",
    "import torch \n",
    "from torch.autograd import Variable\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "filename = \"/home/wafa/Desktop/study/sketchRNN/aaron_sheep.npz\"\n",
    "load_data = np.load(filename, encoding = 'latin1')\n",
    "train_set = load_data['train'][:200]\n",
    "test_set = load_data['test'][:100]\n",
    "\n",
    "nb_steps = 100\n",
    "feature_len=5\n",
    "batch_size = 1\n",
    "max_seq_len=250\n",
    "\n",
    "train_set = utils.DataLoader(train_set, batch_size)\n",
    "test_set = utils.DataLoader(test_set, batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-01-19T22:44:45.785088Z",
     "start_time": "2018-01-19T22:44:45.715420Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, strokeSize, batchSize,  Nhe, Nhd, Nz):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.Nz = Nz\n",
    "        self.Nhe = Nhe\n",
    "        self.Nhd = Nhd\n",
    "        self.cell = nn.LSTM(strokeSize, Nhe//2, 1, bidirectional=True, batch_first=True)\n",
    "        self.mu = nn.Linear(Nhe, Nz)\n",
    "        self.sigma = nn.Linear(Nhe, Nz)\n",
    "        self.h0 = nn.Linear(Nz, Nhd*2) # returns h0 and c0\n",
    "        #print(Nh)\n",
    "\n",
    "    def forward(self, x):\n",
    "        _, (hn, cn) = self.cell(x)\n",
    "        hn = Variable(torch.cat((hn.data[0],hn.data[1]),1))\n",
    "        sigma = self.sigma(hn)\n",
    "        mu = self.mu(hn)\n",
    "        s = sigma\n",
    "        sigma = torch.exp( sigma * 0.5)\n",
    "        eps = Variable(torch.randn(self.Nz))\n",
    "        z = mu + sigma * eps\n",
    "        return F.tanh(self.h0(z)), z, mu, s\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-01-19T23:11:12.957322Z",
     "start_time": "2018-01-19T23:11:12.765455Z"
    }
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, strokeSize, batchSize,  Nhe, Nhd, Nz, Ny):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.Nhe = Nhe\n",
    "        self.Nhd = Nhd\n",
    "        self.Nz = Nz\n",
    "        self.batchSize = batchSize\n",
    "        self.cell = nn.LSTM(strokeSize+Nz, Nhd, 1, batch_first=True)\n",
    "        self.y = nn.Linear(Nhd, Ny)\n",
    "    \n",
    "    def forward(self, x, h0, c0):\n",
    "        output, (hn, cn) = self.cell(x, (h0, c0))\n",
    "        #print(\"avant reshape\", output.size())\n",
    "        #print(output.contiguous().view(-1, self.Nh).size())\n",
    "        output = output.contiguous().view(-1, self.Nhd)\n",
    "        #print(\"after\", output.size())\n",
    "        y = self.y(output)\n",
    "        return y  \n",
    "    \n",
    "    def predict(self, x, h0, c0, z):\n",
    "        out = []\n",
    "        tmp = 1\n",
    "        #h0, z, mu, sigma = self.encoder(x[:, 1:250+1, :])\n",
    "        #new_input = torch.cat((x[:, :250, :], z.view(self.batchSize, 1, self.Nz).expand(self.batchSize, self.max_seq_len, self.Nz)), 2)\n",
    "        #h = h0[:,:self.Nhd]\n",
    "        #c = h0[:,self.Nhd:]\n",
    "        new_input = torch.cat((x[:, 0, :].contiguous().view(self.batchSize, 1, 5), z.view(self.batchSize, 1, self.Nz)), 2)\n",
    "        y, (h, c) = self.cell(new_input, (h0, c0))\n",
    "        output = y.contiguous().view(-1, self.Nhd)\n",
    "        y = self.y(output)\n",
    "        \n",
    "        out.append(y)\n",
    "        z_pen_logits = y[:, -3:]\n",
    "        z_pi, z_mu1, z_mu2, z_sigma1, z_sigma2, z_corr = torch.chunk(y[:, :-3], 6, dim=1)\n",
    "        z_pi = F.softmax(z_pi)\n",
    "        z_pen_logits = F.softmax(z_pen_logits)\n",
    "        z_sigma1 = torch.exp(z_sigma1)\n",
    "        z_sigma2 = torch.exp(z_sigma2)\n",
    "        z_corr = F.tanh(z_corr)\n",
    "        \n",
    "        #normalize the probs to get 1 \n",
    "        probs = np.array(list(z_pi.data[0]))\n",
    "        probs /= probs.sum()\n",
    "        i = np.random.choice(np.arange(0, 2), p= probs)\n",
    "        \n",
    "        mean = [z_mu1.data[0][i], z_mu2.data[0][i]]\n",
    "        cov = [[z_sigma1.data[0][i] * z_sigma1.data[0][i], z_corr.data[0][i] * z_sigma1.data[0][i] * z_sigma2.data[0][i]], [z_corr.data[0][i] * z_sigma1.data[0][i] * z_sigma2.data[0][i], z_sigma2.data[0][i] * z_sigma2.data[0][i]]]\n",
    "        \n",
    "        x = np.random.multivariate_normal(mean, cov, 1)\n",
    "        \n",
    "        probs = np.array(list(z_pen_logits.data[0]))\n",
    "        probs /= probs.sum()\n",
    "        \n",
    "        iPen = np.random.choice(np.arange(0, 3), p=probs)\n",
    "        \n",
    "        pen = [0, 0, 0]\n",
    "        pen[iPen] = 1\n",
    "        \n",
    "        stroke = [x[0][0], x[0][1]] + pen\n",
    "        print(stroke)\n",
    "        \n",
    "        '''new_input = torch.cat((y, z.view(self.batchSize, 1, self.Nz)), 2)\n",
    "        for i in range(1,250):\n",
    "            y, (h, c) = self.cell(new_input, h, c)\n",
    "            out.append(y)\n",
    "        return out\n",
    "        '''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-01-19T23:11:13.414786Z",
     "start_time": "2018-01-19T23:11:13.354435Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class SketchRNN(nn.Module):\n",
    "    \n",
    "    def __init__(self, strokeSize, batchSize, Nhe, Nhd, Nz, Ny, max_seq_len):\n",
    "        super(SketchRNN, self).__init__()\n",
    "        self.batchSize = batchSize\n",
    "        self.Nhe = Nhe\n",
    "        self.Nhd = Nhd\n",
    "        self.max_seq_len = max_seq_len\n",
    "        self.Nz = Nz\n",
    "        self.encoder = Encoder(strokeSize, batchSize, Nhe, Nhd, Nz)\n",
    "        self.decoder = Decoder(strokeSize, batchSize, Nhe, Nhd, Nz, Ny)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # we don't take the inisiale S0 in the encoder so we do the 1:\n",
    "        h0, z, mu, sigma = self.encoder(x[:, 1:250+1, :])\n",
    "        #here we take S0\n",
    "        new_input = torch.cat((x[:, :250, :], z.view(self.batchSize, 1, self.Nz).expand(self.batchSize, self.max_seq_len, self.Nz)), 2)\n",
    "        y = self.decoder(new_input, h0[:,:self.Nhd], h0[:,self.Nhd:])\n",
    "        return y, mu, sigma\n",
    "    \n",
    "    def predict(self, x):\n",
    "        \n",
    "        h0, z, mu, sigma = self.encoder(x[:, 1:250+1, :])\n",
    "        y = self.decoder.predict(x[:, :250, :], h0[:,:self.Nhd], h0[:,self.Nhd:], z)\n",
    "\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-01-19T23:13:17.851096Z",
     "start_time": "2018-01-19T23:13:16.691531Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.021102335091033417, 0.615823257749883, 1, 0, 0]\n",
      "[-0.78079561230249805, 0.014256787788105187, 0, 1, 0]\n",
      "[0.80993878695683108, -0.76205757172356825, 1, 0, 0]\n",
      "[1.1633158694604326, -2.2467039218719513, 0, 1, 0]\n",
      "[2.5294720664366146, -2.352020059199019, 0, 0, 1]\n",
      "[0.2112516582064129, 1.3520768574843387, 0, 1, 0]\n",
      "[0.28230208658823797, 0.1717442292336675, 0, 1, 0]\n",
      "[-0.42966591017080563, 0.77735168226250484, 0, 1, 0]\n",
      "[0.40359272698325688, 0.79893997270539874, 0, 0, 1]\n",
      "[0.78488694934939474, -1.2081477565872407, 0, 0, 1]\n"
     ]
    }
   ],
   "source": [
    "M = 2\n",
    "sketchRnn = SketchRNN(5, 1, 256, 512, 128, 6*M+3, 250)\n",
    "lr=1e-2\n",
    "for _ in range(10):\n",
    "    _, x, s = train_set.random_batch()\n",
    "    x = Variable(torch.from_numpy(x).type(torch.FloatTensor))\n",
    "    sketchRnn.predict(x)\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.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}