notebooks/detr_coco_training_multi_gpu.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Object detection with DETR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade git+https://github.com/EmGarr/kerod.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GPUs configuration from the paper\n",
    "batch_size_per_gpu = 4\n",
    "num_gpus = 16\n",
    "batch_size = batch_size_per_gpu * num_gpus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download and prepare COCO/2017\n",
    "\n",
    "Download and preprocess COCO/2017 to the following format (required by kerod networks):\n",
    "\n",
    "```python\n",
    "dataset = {\n",
    "        'images' : A tensor of float32 and shape [1, height, widht, 3],\n",
    "        'images_info': A tensor of float32 and shape [1, 2] ,\n",
    "        'bbox': A tensor of float32 and shape [1, num_boxes, 4],\n",
    "        'labels': A tensor of int32 and shape [1, num_boxes],\n",
    "        'num_boxes': A tensor of int32 and shape [1, 1],\n",
    "        'weights': A tensor of float32 and shape [1, num_boxes]\n",
    "    }\n",
    "```\n",
    "\n",
    "If you need to download the dataset in a specific directory you can use the argument `data_dir` of `tfds.load`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1cC2k8osNGFw"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "from kerod.core.standard_fields import BoxField, DatasetField\n",
    "from kerod.dataset.preprocessing import preprocess\n",
    "import functools\n",
    "\n",
    "\n",
    "padded_shape = ({\n",
    "  DatasetField.IMAGES: [None, None, 3],\n",
    "  DatasetField.IMAGES_INFO: [2],\n",
    "  DatasetField.IMAGES_PMASK: [None, None]\n",
    "},\n",
    "{\n",
    "  BoxField.BOXES: [None, 4],\n",
    "  BoxField.LABELS: [None],\n",
    "  BoxField.NUM_BOXES: [1],\n",
    "  BoxField.WEIGHTS: [None]\n",
    "})\n",
    "\n",
    "ds_train, ds_info = tfds.load(name=\"coco/2017\", split=\"train\", shuffle_files=True, with_info=True)\n",
    "ds_train = ds_train.map(functools.partial(preprocess, random_crop_size=(384, 600, 3), bgr=True, padded_mask=True),\n",
    "                        num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "# Filter example with no boxes after preprocessing\n",
    "ds_train =  ds_train.filter(lambda x, y: tf.shape(y[BoxField.BOXES])[0] > 1)\n",
    "ds_train =  ds_train.padded_batch(batch_size, padded_shape)\n",
    "ds_train = ds_train.prefetch(tf.data.experimental.AUTOTUNE)\n",
    "\n",
    "ds_val = tfds.load(name=\"coco/2017\", split=\"validation\", shuffle_files=False)\n",
    "ds_val = ds_val.map(functools.partial(preprocess, horizontal_flip=False, bgr=True, padded_mask=True),\n",
    "                      num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "# Filter example with no boxes after preprocessing\n",
    "ds_val =  ds_val.filter(lambda x, y: tf.shape(y[BoxField.BOXES])[0] > 1)\n",
    "ds_val =  ds_val.padded_batch(batch_size, padded_shape)\n",
    "ds_val = ds_val.prefetch(tf.data.experimental.AUTOTUNE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9tt34CM6P-gr"
   },
   "source": [
    "# Load and train the network\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizer selection\n",
    "\n",
    "We train DETR with AdamW setting the initial transformer’s learning rate to **10−4**, the backbone’s to **10−5**, and weight decay to 10−4.\n",
    "We can perform this trick in tensorflow using the arguments **gradient_transformers** of the **optimizer**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def downscale_backbone_grads(grads_and_vars):\n",
    "    \"\"\"Allow to decrease the gradients by 0.1 for the backbone like in the paper.\n",
    "    Decreasing the gradients is equivalent to decreasing the learning rate by 0.1\n",
    "    They observe that having the backbone learning rate roughly an order of magnitude \n",
    "    smaller than the rest of the network is important to stabilize training,\n",
    "    especially in the first few epochs.\n",
    "    \"\"\"\n",
    "    # access to global variable detr defined in the next cells\n",
    "    variables_backbone = {var.name for var in model.backbone.trainable_variables}\n",
    "    return [(0.1 * g if v.name in variables_backbone else g, v)\n",
    "            for g, v in grads_and_vars]\n",
    "\n",
    "def clip_norm_grads(grads_and_vars):\n",
    "    return [(tf.clip_by_norm(g, 0.1), v) for g, v in grads_and_vars]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now just need to call it has follow:\n",
    "\n",
    "```python\n",
    "optimizer = tfa.optimizers.AdamW(weight_decay=1e-4,\n",
    "                                 learning_rate=1e-4,\n",
    "                                 gradient_transformers=[downscale_backbone_grads])\n",
    "```\n",
    "\n",
    "### Run the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from kerod.core.standard_fields import BoxField\n",
    "from kerod.core.learning_rate_schedule import LearningRateScheduler\n",
    "from kerod.model import factory\n",
    "from kerod.model.factory import KerodModel\n",
    "\n",
    "from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint\n",
    "import tensorflow_addons as tfa\n",
    "\n",
    "# Number of classes of COCO\n",
    "classes = ds_info.features['objects']['label'].names\n",
    "num_classes = len(classes)\n",
    "\n",
    "mirrored_strategy = tf.distribute.MirroredStrategy()\n",
    "with mirrored_strategy.scope(): \n",
    "    model = factory.build_model(num_classes, name=KerodModel.detr_resnet50)\n",
    "    learning_rate = 1e-4\n",
    "    optimizer = tfa.optimizers.AdamW(weight_decay=1e-4,\n",
    "                                 learning_rate=learning_rate,\n",
    "                                 gradient_transformers=[downscale_backbone_grads, clip_norm_grads]\n",
    "                                )\n",
    "    model.compile(optimizer=optimizer, loss=None)\n",
    "  \n",
    "callbacks = [\n",
    "    LearningRateScheduler(learning_rate, num_gpus, epochs=[200], use_warmup=False),\n",
    "    TensorBoard('log_tensorboard'),\n",
    "    ModelCheckpoint('checkpoints/')\n",
    "]\n",
    "\n",
    "history = model.fit(ds_train, validation_data=ds_val, epochs=300, callbacks=callbacks)\n",
    "model.save('detr_an_awesome_model')\n",
    "model.save_weights('detr_an_awesome_model.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coco evaluation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow_datasets as tfds\n",
    "\n",
    "ds_val, ds_info = tfds.load(name=\"coco/2017\", split=\"validation\", data_dir='dataset_tensorflow/', shuffle_files=False, with_info=True)\n",
    "# category_ids basicaly map the index 0 the id\n",
    "# e.g: 0 -> 1, 2 -> 3, 79 -> 90\n",
    "category_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Super dirty but the evaluation works\n",
    "\n",
    "#### Perform the analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from kerod.core.standard_fields import DatasetField, BoxField             \n",
    "from kerod.core.box_ops import convert_to_center_coordinates              \n",
    "from kerod.dataset.preprocessing import resize_to_min_dim                 \n",
    "                                                                       \n",
    "results = []                                                           \n",
    "                                                                       \n",
    "for i, example in enumerate(ds_val): \n",
    "    print(i)\n",
    "    # preprocess image \n",
    "    image = example['image'][:, :, ::-1]\n",
    "    image = resize_to_min_dim(image, 800.0, 1333.0)         \n",
    "    image_information = tf.cast(tf.shape(image)[:2], dtype=tf.float32) \n",
    "    inputs = {\n",
    "        'images': tf.expand_dims(image, axis=0),\n",
    "        'images_information':tf.expand_dims(image_information, axis=0)\n",
    "    }\n",
    "                                                                 \n",
    "    # predict                                                          \n",
    "    boxes, scores, labels = model.predict_on_batch(inputs)\n",
    "                                                                       \n",
    "    # Post processing and append to coco results                       \n",
    "    bbox = nmsed_boxes[0] * tf.tile(tf.expand_dims(tf.cast(example['image'].shape[:2], tf.float32)                    \n",
    "        , axis=0), [1, 2])                   \n",
    "    scores = nmsed_scores[0]                                           \n",
    "    labels = nmsed_labels[0]                                           \n",
    "    for i in range(len(scores)):\n",
    "        # scores are sorted so we can remove \n",
    "        if scores[i] < 0.05:\n",
    "            break\n",
    "        sbox = bbox[i].numpy()\n",
    "        # Convert from [y_min, x_min, y_max, x_max] to coco format [x_min, y_min, w, h]\n",
    "        sbox = [sbox[1], sbox[0], sbox[3] - sbox[1], sbox[2] - sbox[0]]\n",
    "        res = {                                                        \n",
    "                'image_id': int(example['image/id']),                       \n",
    "                'category_id': category_ids[int(labels[i])],           \n",
    "                'bbox': [round(float(c), 4) for c in sbox],\n",
    "                'score': round(float(scores[i]), 4),                     \n",
    "            }                                                          \n",
    "        results.append(res)                                            \n",
    "                                                   \n",
    "                                                                       \n",
    "with open('coco_results.json', 'w') as f:                              \n",
    "    json.dump(results, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Install the coco library to compute the performances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip\n",
    "!unzip annotations_trainval2017.zip\n",
    "!pip install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compute the performances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycocotools.coco import COCO\n",
    "from pycocotools.cocoeval import COCOeval\n",
    "\n",
    "with open('coco_results.json', 'r') as f:                              \n",
    "    results = json.load(f)\n",
    "coco = COCO('./annotations/instances_val2017.json')\n",
    "ret = {}\n",
    "\n",
    "cocoDt = coco.loadRes(results)\n",
    "cocoEval = COCOeval(coco, cocoDt, 'bbox')\n",
    "cocoEval.evaluate()\n",
    "cocoEval.accumulate()\n",
    "cocoEval.summarize()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "detr_coco_training_multigpu.ipynb",
   "provenance": []
  },
  "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}