SL_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import keras\n",
    "import os\n",
    "import time\n",
    "import pickle\n",
    "\n",
    "from sl_model import SL512, DSODSL512\n",
    "from ssd_data import InputGenerator\n",
    "from sl_utils import PriorUtil\n",
    "from sl_training import SegLinkLoss, SegLinkFocalLoss\n",
    "\n",
    "from utils.training import Logger, LearningRateDecay\n",
    "from utils.model import load_weights, calc_memory_usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from data_synthtext import GTUtility\n",
    "\n",
    "file_name = 'gt_util_synthtext_seglink.pkl'\n",
    "with open(file_name, 'rb') as f:\n",
    "    gt_util = pickle.load(f)\n",
    "gt_util_train, gt_util_val = gt_util.split(0.9)\n",
    "\n",
    "print(gt_util_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# SegLink\n",
    "model = SL512()\n",
    "weights_path = './models/ssd512_voc_weights_fixed.hdf5'\n",
    "#weights_path = '~/.keras/models/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5'\n",
    "batch_size = 24\n",
    "experiment = 'sl512_synthtext'\n",
    "#experiment = 'sl512fl_synthtext'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# SegLink + DenseNet\n",
    "model = DSODSL512()\n",
    "#model = DSODSL512(activation='leaky_relu')\n",
    "weights_path = None\n",
    "batch_size = 6\n",
    "experiment = 'dsodsl512_synthtext'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# SegLink + ResNet\n",
    "from ssd_model_resnet import SL512_resnet\n",
    "model = SL512_resnet()\n",
    "weights_path = None\n",
    "batch_size = 10\n",
    "experiment = 'sl512_resnet_synthtext'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "if weights_path is not None:\n",
    "    if weights_path.find('ssd512') > -1:\n",
    "        layer_list = [\n",
    "            'conv1_1', 'conv1_2',\n",
    "            'conv2_1', 'conv2_2',\n",
    "            'conv3_1', 'conv3_2', 'conv3_3',\n",
    "            'conv4_1', 'conv4_2', 'conv4_3',\n",
    "            'conv5_1', 'conv5_2', 'conv5_3',\n",
    "            'fc6', 'fc7',\n",
    "            'conv6_1', 'conv6_2',\n",
    "            'conv7_1', 'conv7_2',\n",
    "            'conv8_1', 'conv8_2',\n",
    "            'conv9_1', 'conv9_2',\n",
    "        ]\n",
    "        freeze = [\n",
    "            'conv1_1', 'conv1_2',\n",
    "            'conv2_1', 'conv2_2',\n",
    "            'conv3_1', 'conv3_2', 'conv3_3',\n",
    "            #'conv4_1', 'conv4_2', 'conv4_3',\n",
    "            #'conv5_1', 'conv5_2', 'conv5_3',\n",
    "        ]\n",
    "        \n",
    "        load_weights(model, weights_path, layer_list)\n",
    "        for layer in model.layers:\n",
    "            layer.trainable = not layer.name in freeze\n",
    "    else:\n",
    "        load_weights(model, weights_path)\n",
    "\n",
    "prior_util = PriorUtil(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "epochs = 100\n",
    "initial_epoch = 0\n",
    "\n",
    "gen_train = InputGenerator(gt_util_train, prior_util, batch_size, model.image_size, augmentation=False)\n",
    "gen_val = InputGenerator(gt_util_val, prior_util, batch_size, model.image_size, augmentation=False)\n",
    "\n",
    "checkdir = './checkpoints/' + time.strftime('%Y%m%d%H%M') + '_' + experiment\n",
    "if not os.path.exists(checkdir):\n",
    "    os.makedirs(checkdir)\n",
    "\n",
    "with open(checkdir+'/source.py','wb') as f:\n",
    "    source = ''.join(['# In[%i]\\n%s\\n\\n' % (i, In[i]) for i in range(len(In))])\n",
    "    f.write(source.encode())\n",
    "\n",
    "#optim = keras.optimizers.SGD(lr=1e-3, momentum=0.9, decay=0, nesterov=True)\n",
    "optim = keras.optimizers.Adam(lr=1e-3, beta_1=0.9, beta_2=0.999, epsilon=0.001, decay=0.0)\n",
    "\n",
    "# weight decay\n",
    "regularizer = keras.regularizers.l2(5e-4) # None if disabled\n",
    "#regularizer = None\n",
    "for l in model.layers:\n",
    "    if l.__class__.__name__.startswith('Conv'):\n",
    "        l.kernel_regularizer = regularizer\n",
    "\n",
    "#loss = SegLinkLoss(lambda_offsets=1.0, lambda_links=1.0, neg_pos_ratio=3.0)\n",
    "loss = SegLinkFocalLoss(lambda_segments=100.0, lambda_offsets=1.0, lambda_links=100.0, gamma_segments=2, gamma_links=2)\n",
    "\n",
    "model.compile(optimizer=optim, loss=loss.compute, metrics=loss.metrics)\n",
    "\n",
    "history = model.fit_generator(\n",
    "        gen_train.generate(), \n",
    "        steps_per_epoch=gen_train.num_batches, \n",
    "        epochs=epochs, \n",
    "        verbose=1, \n",
    "        callbacks=[\n",
    "            keras.callbacks.ModelCheckpoint(checkdir+'/weights.{epoch:03d}.h5', verbose=1, save_weights_only=True),\n",
    "            Logger(checkdir),\n",
    "            #LearningRateDecay()\n",
    "        ], \n",
    "        validation_data=gen_val.generate(), \n",
    "        validation_steps=gen_val.num_batches,\n",
    "        class_weight=None,\n",
    "        max_queue_size=1, \n",
    "        workers=1, \n",
    "        #use_multiprocessing=False, \n",
    "        initial_epoch=initial_epoch, \n",
    "        #pickle_safe=False, # will use threading instead of multiprocessing, which is lighter on memory use but slower\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}