model_vsum-ptr-gan.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "feature = pickle.load(open('Dataset/dataset_ytsum/feature.pkl', 'rb'))\n",
    "summary = pickle.load(open('Dataset/dataset_ytsum/summary.pkl', 'rb'))\n",
    "\n",
    "print(len(summary))\n",
    "print(len(feature), list(feature.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import torch as T\n",
    "import torch.nn as nn\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "class DS_YTSum(Dataset):\n",
    "    def __init__(self, feature, summary):\n",
    "        super(DS_YTSum, self).__init__()\n",
    "        \n",
    "        self.feature = feature\n",
    "        self.summary = summary\n",
    "        self.key = list(self.summary.keys())\n",
    "        \n",
    "        MXL_ENC, MXL_DEC = 0, 0\n",
    "        for k in self.key:\n",
    "            MXL_ENC = max(MXL_ENC, len(self.feature[k]))\n",
    "            MXL_DEC = max(MXL_DEC, len(self.summary[k]))\n",
    "        \n",
    "        MXL_ENC += 4\n",
    "        MXL_DEC += 4\n",
    "        \n",
    "        self.enc = dict()\n",
    "        self.dec = dict()\n",
    "        self.ans = dict()\n",
    "        \n",
    "        self.l_in = dict()\n",
    "        self.l_ot = dict()\n",
    "        self.wgt = dict()\n",
    "        \n",
    "        for k in self.key:\n",
    "            enc = np.zeros((MXL_ENC, 4096), dtype=np.float32)\n",
    "            dec = np.zeros((MXL_DEC, 4096), dtype=np.float32)\n",
    "            ans = np.zeros((MXL_DEC, ), dtype=np.int)\n",
    "            wgt = np.zeros((MXL_DEC, ), dtype=np.float32)\n",
    "            \n",
    "            self.l_in[k], self.l_ot[k] = len(self.feature[k])+1, len(self.summary[k])+1\n",
    "            \n",
    "            for i in range(self.l_in[k]-1):\n",
    "                enc[i] = self.feature[k][i]\n",
    "            for i in range(self.l_in[k]-1, MXL_ENC):\n",
    "                enc[i] = self.feature['END']\n",
    "            \n",
    "            dec[0] = self.feature['START']\n",
    "            for i in range(1, self.l_ot[k]):\n",
    "                dec[i] = self.feature[k][self.summary[k][i-1]-1]\n",
    "            for i in range(self.l_ot[k], MXL_DEC):\n",
    "                dec[i] = self.feature['END']\n",
    "            \n",
    "            for i in range(self.l_ot[k]-1):\n",
    "                ans[i] = self.summary[k][i]-1\n",
    "                wgt[i] = 1\n",
    "            wgt[self.l_ot[k]-1] = 1\n",
    "            for i in range(self.l_ot[k]-1, MXL_DEC):\n",
    "                ans[i] = self.l_in[k]-1\n",
    "            \n",
    "            self.enc[k] = enc\n",
    "            self.dec[k] = dec\n",
    "            self.ans[k] = ans\n",
    "            self.wgt[k] = wgt\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.key)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        k = self.key[idx]\n",
    "        \n",
    "        return self.enc[k], self.dec[k], self.ans[k], self.l_in[k], self.l_ot[k], self.wgt[k]\n",
    "\n",
    "ld = DataLoader(DS_YTSum(feature, summary), batch_size=8, shuffle=True)\n",
    "for inp_enc, inp_dec, ans, inp_lin, inp_lot, wgt in ld:\n",
    "    print(inp_enc.shape, inp_dec.shape)\n",
    "    print(ans.shape)\n",
    "    print(inp_lin.shape, inp_lot.shape)\n",
    "    print(wgt.shape)\n",
    "    \n",
    "    print()\n",
    "    print(inp_dec[0][0])\n",
    "    print(inp_enc[0][inp_lin[0]-1])\n",
    "    \n",
    "    print()\n",
    "    print(ans[0])\n",
    "    print(inp_lin[0], inp_lot[0])\n",
    "    \n",
    "    print()\n",
    "    print(wgt[0])\n",
    "    \n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "class Model_G(nn.Module):\n",
    "    def __init__(self, hid=256, dp=0.2):\n",
    "        super(Model_G, self).__init__()\n",
    "        \n",
    "        self.hid = hid\n",
    "        self.dp = dp\n",
    "        \n",
    "        self.gru_enc = nn.GRU(4096, hid, \n",
    "                              num_layers=2, batch_first=True, dropout=dp, bidirectional=True)\n",
    "        self.gru_dec = nn.GRU(4096, hid*2, \n",
    "                              batch_first=True)\n",
    "        self.att = nn.Parameter(T.FloatTensor(hid*2, hid*2))\n",
    "        \n",
    "        self.init()\n",
    "        \n",
    "        self.START = T.from_numpy(feature['START']).view((1, 1, 4096)).float().cuda()\n",
    "        \n",
    "    def init(self):\n",
    "        stdv = 1/math.sqrt(self.hid*2)\n",
    "        \n",
    "        self.att.data.uniform_(stdv, -stdv)\n",
    "    \n",
    "    def rollout(self, inp_enc, out_enc, h, M=1, policy='greedy'):        \n",
    "        STH = h\n",
    "        \n",
    "        outs = []\n",
    "        probs = []\n",
    "        poss = []\n",
    "        for i in range(M):\n",
    "            out = []\n",
    "            prob = []\n",
    "            pos = []\n",
    "            \n",
    "            for j in range(22):\n",
    "                tmp, h = self.gru_dec(self.START if j==0 else tmp, STH if j==0 else h)\n",
    "                tmp = T.bmm(T.matmul(tmp, self.att), out_enc.transpose(1, 2))\n",
    "                out.append(tmp)\n",
    "                \n",
    "                tmp = nn.functional.log_softmax(tmp, dim=2)\n",
    "                prob.append(tmp)\n",
    "                \n",
    "                pb = T.exp(tmp).data.cpu().numpy()[0][0]\n",
    "                if policy=='greedy':\n",
    "                    p = np.argmax(pb)\n",
    "                else:\n",
    "                    p = np.random.choice([i for i in range(len(pb))], p=pb)\n",
    "                pos.append(p)\n",
    "                \n",
    "                tmp = inp_enc[:, p:p+1, :]\n",
    "            \n",
    "            out = T.cat(out, dim=1)\n",
    "            prob = T.cat(prob, dim=1)\n",
    "            \n",
    "            outs.append(out)\n",
    "            probs.append(prob)\n",
    "            poss.append(pos)\n",
    "            \n",
    "        out = T.cat(outs, dim=0)\n",
    "        prob = T.cat(probs, dim=0)\n",
    "        pos = poss\n",
    "        \n",
    "        return out, prob, pos\n",
    "    \n",
    "    def forward(self, inp_enc,  \n",
    "                inp_dec=None, is_tr=True, is_pg=False, M=8):\n",
    "        batch = inp_enc.shape[0]\n",
    "        \n",
    "        outs = []\n",
    "        probs = []\n",
    "        poss = []\n",
    "        for i in range(batch):\n",
    "            out_enc, h = self.gru_enc(inp_enc[i:i+1])\n",
    "            h = h[-2:].view((1, 1, self.hid*2))\n",
    "            \n",
    "            if is_tr==True:\n",
    "                if is_pg==False:\n",
    "                    out_dec, _ = self.gru_dec(inp_dec[i:i+1], h)\n",
    "                    out = T.bmm(T.matmul(out_dec, self.att), out_enc.transpose(1, 2))\n",
    "                    \n",
    "                    outs.append(out)\n",
    "                \n",
    "                else:\n",
    "                    out, prob, pos = self.rollout(inp_enc[i:i+1], out_enc, h, M=M, policy='sample')\n",
    "                    \n",
    "                    outs.append(out)\n",
    "                    probs.append(prob)\n",
    "                    poss.append(pos)\n",
    "            else:\n",
    "                out, prob, pos = self.rollout(inp_enc[i:i+1], out_enc, h, M=1, policy='greedy')\n",
    "                \n",
    "                outs.append(out)\n",
    "                probs.append(prob)\n",
    "                poss.append(pos)\n",
    "        \n",
    "        out = T.cat(outs, dim=0)\n",
    "        if not len(probs)==0:\n",
    "            prob = T.cat(probs, dim=0)\n",
    "        else:\n",
    "            prob = None\n",
    "        pos = [p for pp in poss for p in pp]\n",
    "        \n",
    "        return out, prob, pos\n",
    "\n",
    "model_g = Model_G().cuda()\n",
    "print(model_g)\n",
    "\n",
    "out, _, _ = model_g(inp_enc.cuda(), inp_dec.cuda(), is_tr=True, is_pg=False)\n",
    "print(out.shape)\n",
    "\n",
    "out, prob, pos = model_g(inp_enc.cuda(), is_tr=False)\n",
    "print(out.shape, prob.shape, len(pos), len(pos[0]))\n",
    "\n",
    "out, prob, pos = model_g(inp_enc.cuda(), is_tr=True, is_pg=True)\n",
    "print(out.shape, prob.shape, len(pos), len(pos[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm_notebook as tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 120\n",
    "\n",
    "loss_ce = nn.CrossEntropyLoss(reduction='none').cuda()\n",
    "optim = T.optim.Adam(model_g.parameters(), lr=0.00008)\n",
    "\n",
    "for e in tqdm(range(EPOCHS)):\n",
    "    ls_ep = 0\n",
    "    \n",
    "    model_g.train()\n",
    "    with tqdm(ld) as TQ:\n",
    "        for inp_enc, inp_dec, ans, inp_lin, inp_lot, wgt in TQ:\n",
    "            out, _, _ = model_g(inp_enc.cuda(), inp_dec.cuda(), is_tr=True, is_pg=False)\n",
    "            \n",
    "            out = out.view((-1, 578))\n",
    "            ans = ans.view((-1, )).cuda()\n",
    "            wgt = wgt.view((-1, )).cuda()\n",
    "            \n",
    "            ls_bh = loss_ce(out, ans)*wgt\n",
    "            ls_bh = ls_bh.sum()/wgt.sum()\n",
    "            \n",
    "            optim.zero_grad()\n",
    "            ls_bh.backward()\n",
    "            optim.step()\n",
    "            \n",
    "            ls_bh = ls_bh.cpu().detach().numpy()\n",
    "            ls_ep += ls_bh\n",
    "            \n",
    "            TQ.set_postfix(ls_bh='%.3f'%(ls_bh))\n",
    "            \n",
    "        ls_ep /= len(TQ)\n",
    "        print('Ep %d: %.4f' % (e+1, ls_ep))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_g.load_state_dict(T.load('Model/only_g.pt'))\n",
    "\n",
    "model_g.eval()\n",
    "for inp_enc, inp_dec, ans, inp_lin, inp_lot, wgt in ld:\n",
    "    _, _, pos = model_g(inp_enc.cuda(), inp_dec.cuda(), is_tr=False)\n",
    "    \n",
    "    ans = ans.numpy()\n",
    "    \n",
    "    print(pos[0])\n",
    "    print(ans[0])\n",
    "    print('----------')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model_D(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Model_D, self).__init__()\n",
    "        \n",
    "        self.cnn = nn.Sequential(*[nn.Conv1d(4096, 256, 3, padding=1), nn.ReLU(), \n",
    "                                   nn.Conv1d(256, 128, 3, padding=1), nn.ReLU(), \n",
    "                                   nn.Conv1d(128, 128, 3), nn.ReLU(), \n",
    "                                   nn.Conv1d(128, 64, 3), nn.ReLU()])\n",
    "        self.fc = nn.Sequential(*[nn.Linear(64, 16), nn.ReLU(), \n",
    "                                  nn.Linear(16, 16), nn.ReLU(), \n",
    "                                  nn.Linear(16, 1), nn.Sigmoid()])\n",
    "        \n",
    "    def forward(self, inp):\n",
    "        inp = inp.transpose(1, 2)\n",
    "        out = self.cnn(inp).view((inp.shape[0], -1))\n",
    "        out = self.fc(out)\n",
    "        \n",
    "        return out\n",
    "\n",
    "model_d = Model_D().cuda()\n",
    "print(model_d)\n",
    "\n",
    "inp = T.rand((8, 5, 4096)).cuda()\n",
    "out = model_d(inp)\n",
    "print(out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Loss_PG(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Loss_PG, self).__init__()\n",
    "        \n",
    "        self.bl = 0\n",
    "        self.bn = 0\n",
    "        \n",
    "    def forward(self, prob, pos, lot, rwd):\n",
    "        batch = prob.shape[0]\n",
    "        \n",
    "        ls = 0\n",
    "        cnt = 0\n",
    "        for i in range(batch):\n",
    "            for j in range(lot[i]):\n",
    "                cnt += 1\n",
    "                ls += -prob[i][j][pos[i][j]]*(rwd[i]-self.bl)\n",
    "        ls /= cnt\n",
    "        \n",
    "        self.bl = (self.bl*self.bn + np.average(rwd))/(self.bn+1)\n",
    "        self.bn += 1\n",
    "        \n",
    "        return ls\n",
    "\n",
    "EPOCHS = 40   \n",
    "    \n",
    "loss_mse = nn.MSELoss().cuda()\n",
    "loss_pg = Loss_PG().cuda()\n",
    "\n",
    "optim_g = T.optim.Adam(model_g.parameters(), lr=0.0001)\n",
    "optim_d = T.optim.Adam(model_d.parameters(), lr=0.0001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_cube(inp, pos, lot):\n",
    "    lot = lot.numpy()\n",
    "    pos = sorted(np.random.choice(pos[:lot-1], 5, replace=True if (lot-1)<5 else False))\n",
    "    cube = T.cat([inp[pos[i]].view((1, 1, -1)) for i in range(5)], dim=1)\n",
    "    \n",
    "    return cube.cuda()\n",
    "\n",
    "for e in tqdm(range(EPOCHS)):\n",
    "    ls_g_ep = 0\n",
    "    ls_d_ep = 0\n",
    "    \n",
    "    model_g.train()\n",
    "    with tqdm(ld) as TQ:\n",
    "        for inp_enc, inp_dec, ans, inp_lin, inp_lot, wgt in TQ:\n",
    "            out, prob, pos = model_g(inp_enc.cuda(), inp_dec.cuda(), is_tr=True, is_pg=True, M=4)\n",
    "            \n",
    "            # update G\n",
    "            cube = T.cat([build_cube(inp_enc[i//4], pos[i], inp_lot[i//4]) for i in range(out.shape[0])], dim=0)\n",
    "            rwd = 1-model_d(cube).data.cpu().numpy().squeeze()\n",
    "            ls_bh = loss_pg(prob, pos, np.repeat(inp_lot.numpy(), 4), rwd)\n",
    "            \n",
    "            optim_g.zero_grad()\n",
    "            ls_bh.backward()\n",
    "            optim_g.step()\n",
    "            \n",
    "            ls_bh = ls_bh.cpu().detach().numpy()\n",
    "            ls_g_ep += ls_bh\n",
    "            \n",
    "            # update D (false-case)\n",
    "            out = model_d(cube)\n",
    "            ls_bh = loss_mse(out, T.zeros((out.shape[0])).cuda())\n",
    "            \n",
    "            optim_d.zero_grad()\n",
    "            ls_bh.backward()\n",
    "            optim_d.step()\n",
    "            \n",
    "            ls_bh = ls_bh.cpu().detach().numpy()\n",
    "            ls_d_ep += ls_bh\n",
    "            \n",
    "            # update D (true-case)\n",
    "            cube = T.cat([build_cube(inp_enc[i], ans.numpy()[i], inp_lot[i]) for i in range(inp_enc.shape[0])], dim=0)\n",
    "            \n",
    "            out = model_d(cube)\n",
    "            ls_bh = loss_mse(out, T.ones((out.shape[0])).cuda())\n",
    "            \n",
    "            optim_d.zero_grad()\n",
    "            ls_bh.backward()\n",
    "            optim_d.step()\n",
    "            \n",
    "            ls_bh = ls_bh.cpu().detach().numpy()\n",
    "            ls_d_ep += ls_bh\n",
    "        \n",
    "        ls_g_ep /= len(TQ)\n",
    "        ls_d_ep /= len(TQ)*2\n",
    "        print('Ep %d: G-%.3f D-%.3f' % (e+1, ls_g_ep, ls_d_ep))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_d.load_state_dict(T.load('Model/only_g.pt'))\n",
    "model_g.load_state_dict(T.load('Model/g+d.pt'))"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}