train_rnn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from RNN import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### use gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\"\n",
    "use_gpu = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### speed up and save GPU memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "torch.backends.cudnn.benchmark = True\n",
    "torch.backends.cudnn.enabled = True\n",
    "torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load dataset, namely feature+location a.k.a. combo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "combo_all = dd.io.load('20classes_combo_padzero_1000features.h5') # utilize pad ground truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "gt_all = dd.io.load('20classes_normal_bboxes_all_sqrtwh_list.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "20"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(gt_all)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "20"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(combo_all)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([325, 1004])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "combo_all[0].size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.1394287 ,  0.48230958,  0.16568395,  0.26340988], dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gt_all[0][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### set hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#n_data = 50                   # 20 data needs two process\n",
    "n_features = combo_all[0].size(1)              # YOLO output plus padding ground truth location or zeros\n",
    "hidden_size = input_size = 1004 # input size is the size of input of RNN\n",
    "time_steps = 10               # a.k.a. N=10 in the paper\n",
    "n_batch = 1                   # number of batch\n",
    "n_layers = 1                  # number of recurrent layers\n",
    "N_samples = 3                 # number of sample locations for each time step\n",
    "n_epochs = 1                 # number of epochs for training\n",
    "sigma = torch.FloatTensor(4).cuda()  # variance for (x,y,w,h)\n",
    "sigma[:] = 0.05\n",
    "learning_rate = 0.0007 #0.0006\n",
    "n_dataset = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### save reward result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#result_reward = np.zeros((n_epochs, n_data // time_steps))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### initialize FC module and RNN module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----initializing RNN-----\n"
     ]
    }
   ],
   "source": [
    "rnn = RNN_module(n_features, input_size, hidden_size, n_batch).cuda() # run the module on GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### initialize FC layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for m in rnn.modules():\n",
    "    if isinstance(m, nn.Linear):\n",
    "        m.weight.data.normal_(0, 0.01)\n",
    "        m.bias.data.zero_()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### set the model in \"training mode\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RNN_module(\n",
       "  (fc): Linear(in_features=1004, out_features=1004)\n",
       "  (dropout): Dropout(p=0.5)\n",
       "  (rnn): LSTM(1004, 1004, batch_first=True)\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rnn.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### training process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 1, the dataset 1\n",
      "1st total reward: -26.5858\n",
      "1st total reward: -35.0358\n",
      "1st total reward: -39.9757\n",
      "1st total reward: -52.8693\n",
      "1st total reward: -53.4694\n",
      "1st total reward: -40.1991\n",
      "1st total reward: -39.7892\n",
      "1st total reward: -42.141\n",
      "1st total reward: -37.4087\n",
      "1st total reward: -42.3414\n",
      "epoch 1, the dataset 2\n",
      "1st total reward: -44.905\n",
      "1st total reward: -42.7034\n",
      "1st total reward: -42.7216\n",
      "1st total reward: -38.1536\n",
      "1st total reward: -33.4569\n",
      "1st total reward: -29.5875\n",
      "epoch 1, the dataset 3\n",
      "1st total reward: -32.4583\n",
      "1st total reward: -28.7524\n",
      "1st total reward: -26.9158\n",
      "epoch 1, the dataset 4\n",
      "1st total reward: -29.0239\n",
      "1st total reward: -30.3001\n",
      "1st total reward: -27.2538\n",
      "1st total reward: -24.8728\n",
      "1st total reward: -31.5005\n",
      "1st total reward: -33.5411\n",
      "1st total reward: -32.5758\n",
      "1st total reward: -34.2153\n",
      "1st total reward: -30.5482\n",
      "1st total reward: -32.6993\n",
      "1st total reward: -35.097\n",
      "1st total reward: -34.741\n",
      "1st total reward: -37.8741\n",
      "1st total reward: -38.654\n",
      "1st total reward: -32.5208\n",
      "1st total reward: -34.1414\n",
      "1st total reward: -30.7708\n",
      "1st total reward: -31.8845\n",
      "1st total reward: -25.9774\n",
      "1st total reward: -27.108\n",
      "1st total reward: -27.8322\n",
      "1st total reward: -28.8917\n",
      "1st total reward: -31.7257\n",
      "1st total reward: -25.7842\n",
      "epoch 1, the dataset 5\n",
      "1st total reward: -24.1401\n",
      "1st total reward: -23.887\n",
      "1st total reward: -21.5543\n",
      "1st total reward: -24.2146\n",
      "1st total reward: -21.9267\n",
      "1st total reward: -20.9515\n",
      "1st total reward: -21.1675\n",
      "1st total reward: -20.917\n",
      "1st total reward: -22.6846\n",
      "1st total reward: -22.9906\n",
      "1st total reward: -24.766\n",
      "epoch 1, the dataset 6\n",
      "1st total reward: -29.4426\n",
      "1st total reward: -27.7233\n",
      "1st total reward: -26.8412\n",
      "1st total reward: -26.1968\n",
      "1st total reward: -28.7945\n",
      "1st total reward: -27.6212\n",
      "1st total reward: -26.8237\n",
      "epoch 1, the dataset 7\n",
      "1st total reward: -32.1213\n",
      "1st total reward: -35.1299\n",
      "epoch 1, the dataset 8\n",
      "1st total reward: -23.5666\n",
      "1st total reward: -23.6891\n",
      "1st total reward: -24.6654\n",
      "1st total reward: -24.8242\n",
      "1st total reward: -21.1894\n",
      "1st total reward: -18.6583\n",
      "1st total reward: -20.2395\n",
      "1st total reward: -18.2089\n",
      "1st total reward: -19.0135\n",
      "1st total reward: -19.9799\n",
      "1st total reward: -18.6647\n",
      "epoch 1, the dataset 9\n",
      "1st total reward: -20.3393\n",
      "1st total reward: -20.9742\n",
      "1st total reward: -24.3652\n",
      "1st total reward: -16.6269\n",
      "1st total reward: -15.8307\n",
      "epoch 1, the dataset 10\n",
      "1st total reward: -29.3461\n",
      "1st total reward: -23.7523\n",
      "1st total reward: -18.8668\n",
      "1st total reward: -18.3814\n",
      "1st total reward: -20.7595\n"
     ]
    }
   ],
   "source": [
    "reward_record = []\n",
    "reward_2nd_record = []\n",
    "for epoch in range(n_epochs): # run the policy for n_epochs\n",
    "    for ith_dataset in range(n_dataset):\n",
    "        n_subtrain = gt_all[ith_dataset].shape[0] // 3 // time_steps\n",
    "        n_subtrain *= time_steps\n",
    "        combo_train = combo_all[ith_dataset][:n_subtrain]\n",
    "        gt_train = gt_all[ith_dataset][:n_subtrain]\n",
    "        #gt = Variable(torch.from_numpy(gt_train).float(), volatile=True) # run this variable on GPU\n",
    "        #combo = Variable(combo_train).cuda() # run this variable on GPU\n",
    "        \n",
    "        loop_begin = 0  # loop pointer\n",
    "        gt_pointer = 0  # ground truth array pointer\n",
    "        #print('')\n",
    "        print('epoch %d, the dataset %d' %(epoch+1, ith_dataset+1))\n",
    "        # randomly generate initial hidden and cell states at begining of each epoch\n",
    "        h_state = Variable(torch.randn(n_layers, n_batch, hidden_size)).cuda()\n",
    "        c_state = Variable(torch.randn(n_layers, n_batch, hidden_size)).cuda()\n",
    "        for loop in range(n_subtrain // time_steps): # for each 10 frames\n",
    "            #### forward pass ####\n",
    "            # get input of RNN from FC layer output\n",
    "            input = combo_train[loop_begin:loop_begin + time_steps]\n",
    "            input = Variable(input).cuda()\n",
    "            #print(input.size())\n",
    "            \n",
    "            loop_begin += time_steps\n",
    "\n",
    "            # compute mu from hidden states and sample N_samples location for each time step t\n",
    "            sample_location = Variable(torch.FloatTensor(n_batch, time_steps, N_samples, 4), volatile=True).cuda()  # 4 is location (x,y,w,h)\n",
    "            mu_tensor = Variable(torch.FloatTensor(time_steps, 4)).cuda()\n",
    "            for index_batch in range(n_batch):  # index for batch\n",
    "                for t in range(time_steps):  # index for time_step\n",
    "                    # for each time step apply RNN and get h_state, c_state for current time step\n",
    "                    #one_input = input[t,:].view(1,-1)\n",
    "                    #one_input = Variable(one_input).cuda()\n",
    "                    h_state, c_state = rnn(input[t,:].view(1,-1), h_state, c_state)\n",
    "                    # compute network output mean mu of location which contains (x,y,w,h)\n",
    "                    mu = h_state[0, 0][-4:]  # [0,0] means the first recurrent layer and batch\n",
    "                    mu_tensor[t] = mu\n",
    "                    # randomly sample N location predictions for current time step\n",
    "                    m = Normal(mu, sigma)\n",
    "                    for index_sample in range(N_samples):  # sample n_samples from Gaussian distribution with mean mu and var sigma\n",
    "                        # randomly sample predictions for N episodes, namely N_samples\n",
    "                        sample_location[index_batch, t, index_sample, :] = m.sample().clone()\n",
    "                    # delete mu and m at the end of each loop in order to save memory\n",
    "\n",
    "            # compute the first reward\n",
    "            one_gt = gt_train[gt_pointer:gt_pointer + time_steps]\n",
    "            gt = Variable(torch.from_numpy(one_gt).float(), volatile=True)\n",
    "            gt = gt.cuda()\n",
    "            rep_gt = gt.repeat(1, N_samples).view(1, time_steps, N_samples, 4)\n",
    "            gt_pointer += time_steps\n",
    "            abs_sub = (sample_location - rep_gt).abs()\n",
    "            avg_val = abs_sub.mean(dim=-1)\n",
    "            max_val = abs_sub.max(dim=-1)[0]\n",
    "            reward = -avg_val - max_val\n",
    "            G_display = reward.sum().cpu().data.numpy()[0]\n",
    "            reward_record.append(G_display)\n",
    "            print('1st total reward:', G_display)\n",
    "            \n",
    "            # compute baseline\n",
    "            baseline = reward.sum(dim=-1) / time_steps\n",
    "            #baseline = Variable(baseline.data, volatile=True).cuda()\n",
    "            baseline_rep = baseline.view(-1,1).repeat(1,N_samples)\n",
    "            R_b = (reward - baseline_rep) # (R^t_i - b_t) with shape=(N*T)\n",
    "            \n",
    "            '''\n",
    "            # compute the second reward\n",
    "            #box1_gt = gt_train[gt_pointer:gt_pointer + time_steps]  # (10, 4)\n",
    "            reward_2 = Variable(torch.zeros(time_steps, N_samples)).cuda()\n",
    "            for k in range(N_samples):\n",
    "                box2_pred = sample_location.squeeze()[:, k, :]\n",
    "                iou = compute_iou(gt, box2_pred)\n",
    "                reward_2[:, k] = iou.diag()\n",
    "            # print(reward_2.size())\n",
    "            G = reward_2.sum()  # total reward or the expectation of total reward\n",
    "            total_reward_2 = G.cpu().data.numpy()[0]\n",
    "            if total_reward_2 != 0:\n",
    "                print('2nd total reward:', total_reward_2)      \n",
    "                \n",
    "            if epoch < 10 or total_reward_2 == 0:\n",
    "                # compute baseline\n",
    "                baseline = reward.sum(dim=-1) / time_steps\n",
    "                #baseline = Variable(baseline.data, volatile=True).cuda()\n",
    "                baseline_rep = baseline.view(-1,1).repeat(1,N_samples)\n",
    "                R_b = (reward - baseline_rep) # (R^t_i - b_t) with shape=(N*T)\n",
    "            else:\n",
    "                # compute baseline\n",
    "                baseline = reward_2.sum(dim=-1) / time_steps\n",
    "                #baseline = Variable(baseline.data, volatile=True).cuda()\n",
    "                baseline_rep = baseline.view(-1,1).repeat(1,N_samples)\n",
    "                R_b = (reward_2 - baseline_rep) # (R^t_i - b_t) with shape=(N*T)            \n",
    "            '''\n",
    "\n",
    "            \n",
    "            \n",
    "            #### backward pass ####\n",
    "            # get parameter size list of all parameter tensor\n",
    "            size_list = []\n",
    "            for param in rnn.parameters():\n",
    "                size_list.append(list(param.size()))\n",
    "\n",
    "            # compute dimension of each parameter vector\n",
    "            param_size = 0\n",
    "            truncate_size = []  # store size of each parameters\n",
    "            for size in size_list:\n",
    "                if len(size) == 2:\n",
    "                    param_size = size[0] * size[1] + param_size\n",
    "                    truncate_size.append(size[0] * size[1])\n",
    "                else:\n",
    "                    param_size += size[0]\n",
    "                    truncate_size.append(size[0])\n",
    "            truncate_size = np.array(truncate_size)\n",
    "\n",
    "            # compute gradient of mu w.r.t. W\n",
    "            gradient_mu = torch.FloatTensor(time_steps, 4, param_size)  ## it need to add batch_index as well\n",
    "            for t in range(mu_tensor.size(0)):  # time_step\n",
    "                for l in range(mu_tensor.size(1)):  # 4-dim location which contains (x,y,w,h)\n",
    "                    mu_tensor[t, l].backward(retain_graph=True)\n",
    "                    # compute grad. of fc layer\n",
    "                    for index, param in enumerate(rnn.parameters()):\n",
    "                        if index == 0:\n",
    "                            param_vector = param.grad.data.clone().view(-1)\n",
    "                        else:\n",
    "                            param_vector = torch.cat((param_vector, param.grad.data.clone().view(-1)))\n",
    "                        param.grad.data.zero_()\n",
    "                    gradient_mu[t, l, :] = param_vector\n",
    "\n",
    "            # compute gradient of policy w.r.t. mu\n",
    "            difference = sample_location - mu_tensor.repeat(1, N_samples).view(1, time_steps, N_samples, 4)\n",
    "            gradient_policy = difference / sigma[0]**2\n",
    "\n",
    "            # compute gradient of W using chain rule\n",
    "            # [0,t] means batch index 0 at time step t\n",
    "            gradient_W = torch.bmm(gradient_policy.data.squeeze(), gradient_mu.cuda())\n",
    "            factor1 = R_b.squeeze().view(time_steps,N_samples,1)\n",
    "            factor2 = factor1.repeat(1,1,param_size)#.cuda()\n",
    "            factor = factor2.view(time_steps, N_samples,param_size)\n",
    "            gradient_G = torch.mul(gradient_W, factor.data)\n",
    "\n",
    "            # sum up gradient of W for N_sample at all time_steps\n",
    "            gradient_sum = gradient_G.sum(0).sum(0) / N_samples\n",
    "\n",
    "\n",
    "            #### update parameters ####\n",
    "            # update parameters of fc model using gradient ascent\n",
    "            pointer = 0  # indicate the i-th parameter\n",
    "            begin = 0  # indicate beginning position\n",
    "            for param in rnn.parameters():\n",
    "                param.data += learning_rate * gradient_sum[begin:begin + truncate_size[pointer]].cuda().view(size_list[pointer])\n",
    "                begin += truncate_size[pointer]\n",
    "                pointer += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### set the model in \"testing mode\" in order to close Dropout and save the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RNN_module(\n",
       "  (fc): Linear(in_features=1004, out_features=1004)\n",
       "  (dropout): Dropout(p=0.5)\n",
       "  (rnn): LSTM(1004, 1004, batch_first=True)\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rnn.eval()  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### save model and reward record"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "reward_record = np.array(reward_record)\n",
    "dd.io.save('./saved_models/reward_04110942', reward_record)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "torch.save(rnn.state_dict(),'./saved_models/DRLT_10classes_1000features_40iters_04110942.pth')"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}