SSD_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import keras\n",
    "import os, time\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from ssd_model import SSD300, SSD512\n",
    "from ssd_utils import PriorUtil\n",
    "from ssd_data import InputGenerator\n",
    "from ssd_training import SSDLoss, SSDFocalLoss\n",
    "\n",
    "from utils.model import load_weights\n",
    "from utils.training import MetricUtility"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pascal VOC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from data_voc import GTUtility\n",
    "gt_util = GTUtility('data/VOC2007/')\n",
    "gt_util_train, gt_util_val = gt_util.split(0.8)\n",
    "experiment = 'ssd512_voc'\n",
    "#print(gt_util)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### MS COCO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from data_coco import GTUtility\n",
    "gt_util = GTUtility('data/COCO/')\n",
    "gt_util_train, gt_util_val = gt_util.split(0.8)\n",
    "experiment = 'ssd512_coco'\n",
    "#print(gt_util)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model = SSD300(num_classes=gt_util.num_classes)\n",
    "model = SSD512(num_classes=gt_util.num_classes)\n",
    "\n",
    "prior_util = PriorUtil(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_epoch = 0\n",
    "\n",
    "#!wget -O ./models/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5 https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5\n",
    "weights_path = './models/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5'\n",
    "layer_list = [('block1_conv1', 'conv1_1'),\n",
    "              ('block1_conv2', 'conv1_2'),\n",
    "              ('block2_conv1', 'conv2_1'),\n",
    "              ('block2_conv2', 'conv2_2'),\n",
    "              ('block3_conv1', 'conv3_1'),\n",
    "              ('block3_conv2', 'conv3_2'),\n",
    "              ('block3_conv3', 'conv3_3'),\n",
    "              ('block4_conv1', 'conv4_1'),\n",
    "              ('block4_conv2', 'conv4_2'),\n",
    "              ('block4_conv3', 'conv4_3'),\n",
    "              ('block5_conv1', 'conv5_1'),\n",
    "              ('block5_conv2', 'conv5_2'),\n",
    "              ('block5_conv3', 'conv5_3')]\n",
    "#load_weights(model, weights_path, layer_list)\n",
    "\n",
    "weights_path = './models/ssd512_voc_weights_fixed.hdf5'\n",
    "#weights_path = './checkpoints/201709251316_ssd512_coco/weights.040.h5'; initial_epoch = 41\n",
    "#weights_path = './checkpoints/201709081909_ssd512_voc/weights.034.h5'; initial_epoch = 35\n",
    "#model.load_weights(weights_path, by_name=True)\n",
    "load_weights(model, weights_path)\n",
    "\n",
    "freeze = ['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",
    "         ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 100\n",
    "batch_size = 24\n",
    "\n",
    "#optimizer = tf.optimizers.SGD(learning_rate=1e-3, momentum=0.9, decay=0, nesterov=True)\n",
    "optimizer = tf.optimizers.Adam(learning_rate=1e-4, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)\n",
    "\n",
    "#loss = SSDLoss(alpha=1.0, neg_pos_ratio=3.0)\n",
    "loss = SSDFocalLoss()\n",
    "\n",
    "#regularizer = None\n",
    "regularizer = keras.regularizers.l2(5e-4) # None if disabled\n",
    "\n",
    "gen_train = InputGenerator(gt_util_train, prior_util, batch_size, model.image_size, augmentation=True)\n",
    "gen_val = InputGenerator(gt_util_val, prior_util, batch_size, model.image_size, augmentation=True)\n",
    "\n",
    "\n",
    "dataset_train, dataset_val = gen_train.get_dataset(), gen_val.get_dataset()\n",
    "iterator_train, iterator_val = iter(dataset_train), iter(dataset_val)\n",
    "\n",
    "checkdir = './checkpoints/' + time.strftime('%Y%m%d%H%M') + '_' + experiment\n",
    "\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",
    "print(checkdir)\n",
    "\n",
    "for l in model.layers:\n",
    "    l.trainable = not l.name in freeze\n",
    "    if regularizer and l.__class__.__name__.startswith('Conv'):\n",
    "        model.add_loss(lambda l=l: regularizer(l.kernel))\n",
    "\n",
    "metric_util = MetricUtility(loss.metric_names, logdir=checkdir)\n",
    "\n",
    "@tf.function\n",
    "def step(x, y_true, training=False):\n",
    "    if training:\n",
    "        with tf.GradientTape() as tape:\n",
    "            y_pred = model(x, training=True)\n",
    "            metric_values = loss.compute(y_true, y_pred)\n",
    "            total_loss = metric_values['loss']\n",
    "            if len(model.losses):\n",
    "                total_loss += tf.add_n(model.losses)\n",
    "        gradients = tape.gradient(total_loss, model.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "    else:\n",
    "        y_pred = model(x, training=True)\n",
    "        metric_values = loss.compute(y_true, y_pred)\n",
    "    return metric_values\n",
    "\n",
    "for k in tqdm(range(initial_epoch, epochs), 'total', leave=False):\n",
    "    print('\\nepoch %i/%i' % (k+1, epochs))\n",
    "    metric_util.on_epoch_begin()\n",
    "\n",
    "    for i in tqdm(range(gen_train.num_batches), 'training', leave=False):\n",
    "        x, y_true = next(iterator_train)\n",
    "        metric_values = step(x, y_true, training=True)\n",
    "        metric_util.update(metric_values, training=True)\n",
    "    \n",
    "    model.save_weights(checkdir+'/weights.%03i.h5' % (k+1,))\n",
    "\n",
    "    for i in tqdm(range(gen_val.num_batches), 'validation', leave=False):\n",
    "        x, y_true = next(iterator_val)\n",
    "        metric_values = step(x, y_true, training=False)\n",
    "        metric_util.update(metric_values, training=False)\n",
    "\n",
    "    metric_util.on_epoch_end(verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.8.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}