{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# D2Go Beginner's Tutorial\n",
    "\n",
    "This is beginner tutorial for [d2go project](https://github.com/facebookresearch/d2go). We will go through some basic usage of d2go, including:\n",
    "- Run inference on images or videos, with a pretrained d2go model\n",
    "- Load a new dataset and train a d2go model\n",
    "- Export models to int8 using post-training quantization. \n",
    "\n",
    "Please install d2go before running this tutorial following the [instructions](https://github.com/facebookresearch/d2go)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference with Pre-trained Models\n",
    "\n",
    "In this section, we will show how to load pretrained models using d2go model_zoo API, and how to make predictions with d2go models and visualize the output. \n",
    "\n",
    "- First import the model zoo API from d2go and get a pretrained Faster R-CNN model with FBNetV3 backbone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "** fvcore version of PathManager will be deprecated soon. **\n",
      "** Please migrate to the version in iopath repo. **\n",
      "https://github.com/facebookresearch/iopath \n",
      "\n",
      "Arguments ['width_divisor', 'dw_skip_bnrelu', 'zero_last_bn_gamma'] skipped for op Conv2d\n"
     ]
    }
   ],
   "source": [
    "from d2go.model_zoo import model_zoo\n",
    "model = model_zoo.get('faster_rcnn_fbnetv3a_C4.yaml', trained=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Download an image from the COCO dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f3b75c345b0>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 216108 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cv2\n",
    "from matplotlib import pyplot as plt\n",
    "!wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg\n",
    "im = cv2.imread(\"./input.jpg\")\n",
    "plt.imshow(im)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Then we can create a `DemoPredictor` to run inference on this image and see the raw outputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([17,  0, 19, 17, 17, 17, 19,  0,  0, 19, 17, 17, 19], device='cuda:0')\n",
      "Boxes(tensor([[139.2618, 238.0369, 467.4290, 480.0000],\n",
      "        [271.1126, 150.1942, 341.3032, 406.9402],\n",
      "        [505.8431, 268.7002, 585.4443, 339.0620],\n",
      "        [  2.2081, 282.8764,  78.0862, 475.8036],\n",
      "        [336.2187, 241.6059, 424.0182, 301.4847],\n",
      "        [262.5104, 226.0209, 468.7945, 413.3305],\n",
      "        [ 46.3860, 282.2225,  74.5961, 344.6837],\n",
      "        [  0.0000, 285.9752,  75.4428, 477.4194],\n",
      "        [ 46.6597, 281.7245,  74.5386, 344.4575],\n",
      "        [  0.0000, 286.2906,  75.8711, 477.1982],\n",
      "        [505.0901, 262.1472, 586.4306, 323.4085],\n",
      "        [ 46.7585, 282.4321,  74.5174, 344.7928],\n",
      "        [559.0256, 275.1691, 595.4606, 339.6431]], device='cuda:0'))\n"
     ]
    }
   ],
   "source": [
    "from d2go.utils.demo_predictor import DemoPredictor\n",
    "predictor = DemoPredictor(model)\n",
    "outputs = predictor(im)\n",
    "# the output object categories and corresponding bounding boxes\n",
    "print(outputs[\"instances\"].pred_classes)\n",
    "print(outputs[\"instances\"].pred_boxes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Let's visualize the output predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f3b7420c3a0>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 216476 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "from detectron2.utils.visualizer import Visualizer\n",
    "from detectron2.data import MetadataCatalog, DatasetCatalog\n",
    "\n",
    "v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(\"coco_2017_train\"))\n",
    "out = v.draw_instance_predictions(outputs[\"instances\"].to(\"cpu\"))\n",
    "plt.imshow(out.get_image()[:, :, ::-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train on a custom dataset\n",
    "In this section, we show how to train a d2go model on a custom dataset.\n",
    "\n",
    "We use [the balloon segmentation dataset](https://github.com/matterport/Mask_RCNN/tree/master/samples/balloon)\n",
    "which only has one class: balloon.\n",
    "We'll train a balloon segmentation model from an existing model pre-trained on COCO dataset, available in d2go's [model zoo](https://github.com/facebookresearch/d2go/blob/master/MODEL_ZOO.md).\n",
    "\n",
    "### Prepare the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2021-03-03 22:12:03--  https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip\n",
      "Resolving github.com (github.com)... 192.30.255.112\n",
      "Connecting to github.com (github.com)|192.30.255.112|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://github-releases.githubusercontent.com/107595270/737339e2-2b83-11e8-856a-188034eb3468?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210303%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210303T221203Z&X-Amz-Expires=300&X-Amz-Signature=7390c6892812465f3b14733b7fc8a29ce55248b8c99e2ea94694b6ed1fda1ec2&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=107595270&response-content-disposition=attachment%3B%20filename%3Dballoon_dataset.zip&response-content-type=application%2Foctet-stream [following]\n",
      "--2021-03-03 22:12:03--  https://github-releases.githubusercontent.com/107595270/737339e2-2b83-11e8-856a-188034eb3468?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210303%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210303T221203Z&X-Amz-Expires=300&X-Amz-Signature=7390c6892812465f3b14733b7fc8a29ce55248b8c99e2ea94694b6ed1fda1ec2&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=107595270&response-content-disposition=attachment%3B%20filename%3Dballoon_dataset.zip&response-content-type=application%2Foctet-stream\n",
      "Resolving github-releases.githubusercontent.com (github-releases.githubusercontent.com)... 185.199.109.154, 185.199.110.154, 185.199.111.154, ...\n",
      "Connecting to github-releases.githubusercontent.com (github-releases.githubusercontent.com)|185.199.109.154|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 38741381 (37M) [application/octet-stream]\n",
      "Saving to: ‘balloon_dataset.zip.5’\n",
      "\n",
      "balloon_dataset.zip 100%[===================>]  36.95M  58.3MB/s    in 0.6s    \n",
      "\n",
      "2021-03-03 22:12:04 (58.3 MB/s) - ‘balloon_dataset.zip.5’ saved [38741381/38741381]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# download, decompress the data\n",
    "!wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip\n",
    "!unzip -o balloon_dataset.zip > /dev/null"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "D2Go is built on top of detectron2. Let's register the balloon dataset to detectron2, following the [detectron2 custom dataset tutorial](https://detectron2.readthedocs.io/tutorials/datasets.html).\n",
    "Here, the dataset is in its custom format, therefore we write a function to parse it and prepare it into detectron2's standard format. User should write such a function when using a dataset in custom format. See the tutorial for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if your dataset is in COCO format, this cell can be replaced by the following three lines:\n",
    "# from detectron2.data.datasets import register_coco_instances\n",
    "# register_coco_instances(\"my_dataset_train\", {}, \"json_annotation_train.json\", \"path/to/image/dir\")\n",
    "# register_coco_instances(\"my_dataset_val\", {}, \"json_annotation_val.json\", \"path/to/image/dir\")\n",
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "from detectron2.structures import BoxMode\n",
    "\n",
    "def get_balloon_dicts(img_dir):\n",
    "    json_file = os.path.join(img_dir, \"via_region_data.json\")\n",
    "    with open(json_file) as f:\n",
    "        imgs_anns = json.load(f)\n",
    "\n",
    "    dataset_dicts = []\n",
    "    for idx, v in enumerate(imgs_anns.values()):\n",
    "        record = {}\n",
    "        \n",
    "        filename = os.path.join(img_dir, v[\"filename\"])\n",
    "        height, width = cv2.imread(filename).shape[:2]\n",
    "        \n",
    "        record[\"file_name\"] = filename\n",
    "        record[\"image_id\"] = idx\n",
    "        record[\"height\"] = height\n",
    "        record[\"width\"] = width\n",
    "      \n",
    "        annos = v[\"regions\"]\n",
    "        objs = []\n",
    "        for _, anno in annos.items():\n",
    "            assert not anno[\"region_attributes\"]\n",
    "            anno = anno[\"shape_attributes\"]\n",
    "            px = anno[\"all_points_x\"]\n",
    "            py = anno[\"all_points_y\"]\n",
    "            poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]\n",
    "            poly = [p for x in poly for p in x]\n",
    "\n",
    "            obj = {\n",
    "                \"bbox\": [np.min(px), np.min(py), np.max(px), np.max(py)],\n",
    "                \"bbox_mode\": BoxMode.XYXY_ABS,\n",
    "                \"segmentation\": [poly],\n",
    "                \"category_id\": 0,\n",
    "            }\n",
    "            objs.append(obj)\n",
    "        record[\"annotations\"] = objs\n",
    "        dataset_dicts.append(record)\n",
    "    return dataset_dicts\n",
    "\n",
    "for d in [\"train\", \"val\"]:\n",
    "    DatasetCatalog.register(\"balloon_\" + d, lambda d=d: get_balloon_dicts(\"balloon/\" + d))\n",
    "    MetadataCatalog.get(\"balloon_\" + d).set(thing_classes=[\"balloon\"], evaluator_type=\"coco\")\n",
    "balloon_metadata = MetadataCatalog.get(\"balloon_train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To verify the data loading is correct, let's visualize the annotations of randomly selected samples in the training set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48912 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 155084 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 78496 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "dataset_dicts = get_balloon_dicts(\"balloon/train\")\n",
    "for d in random.sample(dataset_dicts, 3):\n",
    "    img = cv2.imread(d[\"file_name\"])\n",
    "    visualizer = Visualizer(img[:, :, ::-1], metadata=balloon_metadata, scale=0.5)\n",
    "    out = visualizer.draw_dataset_dict(d)\n",
    "    plt.figure()\n",
    "    plt.imshow(out.get_image()[:, :, ::-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train\n",
    "Now, let's fine-tune a COCO-pretrained FBNetV3A Mask R-CNN model on the balloon dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "for d in [\"train\", \"val\"]:\n",
    "    MetadataCatalog.get(\"balloon_\" + d).set(thing_classes=[\"balloon\"], evaluator_type=\"coco\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Arguments ['width_divisor', 'dw_skip_bnrelu', 'zero_last_bn_gamma'] skipped for op Conv2d\n",
      "Skip loading parameter 'roi_heads.box_predictor.cls_score.weight' to the model due to incompatible shapes: (81, 200) in the checkpoint but (2, 200) in the model! You might want to double check if this is expected.\n",
      "Skip loading parameter 'roi_heads.box_predictor.cls_score.bias' to the model due to incompatible shapes: (81,) in the checkpoint but (2,) in the model! You might want to double check if this is expected.\n",
      "Skip loading parameter 'roi_heads.box_predictor.bbox_pred.weight' to the model due to incompatible shapes: (320, 200) in the checkpoint but (4, 200) in the model! You might want to double check if this is expected.\n",
      "Skip loading parameter 'roi_heads.box_predictor.bbox_pred.bias' to the model due to incompatible shapes: (320,) in the checkpoint but (4,) in the model! You might want to double check if this is expected.\n",
      "Using previously cached COCO format annotations at './output/inference/default/600/balloon_val/balloon_val_coco_format.json'. You need to clear the cache file if your dataset has been modified.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading and preparing results...\n",
      "DONE (t=0.00s)\n",
      "creating index...\n",
      "index created!\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.494\n",
      " Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.651\n",
      " Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.543\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.104\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.757\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.204\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.526\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.526\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.118\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.810\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'model_final': CfgNode({'VERSION': 2, 'MODEL': CfgNode({'LOAD_PROPOSALS': False, 'MASK_ON': False, 'KEYPOINT_ON': False, 'DEVICE': 'cuda', 'META_ARCHITECTURE': 'GeneralizedRCNN', 'WEIGHTS': './output/model_final.pth', 'PIXEL_MEAN': [103.53, 116.28, 123.675], 'PIXEL_STD': [1.0, 1.0, 1.0], 'BACKBONE': CfgNode({'NAME': 'FBNetV2C4Backbone', 'FREEZE_AT': 2}), 'FPN': CfgNode({'IN_FEATURES': [], 'OUT_CHANNELS': 256, 'NORM': '', 'FUSE_TYPE': 'sum'}), 'PROPOSAL_GENERATOR': CfgNode({'NAME': 'RPN', 'MIN_SIZE': 0}), 'ANCHOR_GENERATOR': CfgNode({'NAME': 'DefaultAnchorGenerator', 'SIZES': [[32, 64, 96, 128, 160]], 'ASPECT_RATIOS': [[0.5, 1.0, 2.0]], 'ANGLES': [[-90, 0, 90]], 'OFFSET': 0.0}), 'RPN': CfgNode({'HEAD_NAME': 'FBNetV2RpnHead', 'IN_FEATURES': ['trunk3'], 'BOUNDARY_THRESH': -1, 'IOU_THRESHOLDS': [0.3, 0.7], 'IOU_LABELS': [0, -1, 1], 'BATCH_SIZE_PER_IMAGE': 256, 'POSITIVE_FRACTION': 0.5, 'BBOX_REG_LOSS_TYPE': 'smooth_l1', 'BBOX_REG_LOSS_WEIGHT': 1.0, 'BBOX_REG_WEIGHTS': (1.0, 1.0, 1.0, 1.0), 'SMOOTH_L1_BETA': 0.0, 'LOSS_WEIGHT': 1.0, 'PRE_NMS_TOPK_TRAIN': 2000, 'PRE_NMS_TOPK_TEST': 1000, 'POST_NMS_TOPK_TRAIN': 2000, 'POST_NMS_TOPK_TEST': 30, 'NMS_THRESH': 0.7}), 'ROI_HEADS': CfgNode({'NAME': 'StandardROIHeads', 'NUM_CLASSES': 1, 'IN_FEATURES': ['trunk3'], 'IOU_THRESHOLDS': [0.5], 'IOU_LABELS': [0, 1], 'BATCH_SIZE_PER_IMAGE': 128, 'POSITIVE_FRACTION': 0.25, 'SCORE_THRESH_TEST': 0.05, 'NMS_THRESH_TEST': 0.5, 'PROPOSAL_APPEND_GT': True}), 'ROI_BOX_HEAD': CfgNode({'NAME': 'FBNetV2RoIBoxHead', 'BBOX_REG_LOSS_TYPE': 'smooth_l1', 'BBOX_REG_LOSS_WEIGHT': 1.0, 'BBOX_REG_WEIGHTS': (10.0, 10.0, 5.0, 5.0), 'SMOOTH_L1_BETA': 0.0, 'POOLER_RESOLUTION': 6, 'POOLER_SAMPLING_RATIO': 0, 'POOLER_TYPE': 'ROIAlignV2', 'NUM_FC': 0, 'FC_DIM': 1024, 'NUM_CONV': 0, 'CONV_DIM': 256, 'NORM': 'naiveSyncBN', 'CLS_AGNOSTIC_BBOX_REG': False, 'TRAIN_ON_PRED_BOXES': False}), 'ROI_BOX_CASCADE_HEAD': CfgNode({'BBOX_REG_WEIGHTS': ((10.0, 10.0, 5.0, 5.0), (20.0, 20.0, 10.0, 10.0), (30.0, 30.0, 15.0, 15.0)), 'IOUS': (0.5, 0.6, 0.7)}), 'ROI_MASK_HEAD': CfgNode({'NAME': 'MaskRCNNConvUpsampleHead', 'POOLER_RESOLUTION': 14, 'POOLER_SAMPLING_RATIO': 0, 'NUM_CONV': 4, 'CONV_DIM': 256, 'NORM': '', 'CLS_AGNOSTIC_MASK': False, 'POOLER_TYPE': 'ROIAlignV2'}), 'ROI_KEYPOINT_HEAD': CfgNode({'NAME': 'KRCNNConvDeconvUpsampleHead', 'POOLER_RESOLUTION': 14, 'POOLER_SAMPLING_RATIO': 0, 'CONV_DIMS': (512, 512, 512, 512, 512, 512, 512, 512), 'NUM_KEYPOINTS': 17, 'MIN_KEYPOINTS_PER_IMAGE': 1, 'NORMALIZE_LOSS_BY_VISIBLE_KEYPOINTS': True, 'LOSS_WEIGHT': 1.0, 'POOLER_TYPE': 'ROIAlignV2'}), 'SEM_SEG_HEAD': CfgNode({'NAME': 'SemSegFPNHead', 'IN_FEATURES': ['p2', 'p3', 'p4', 'p5'], 'IGNORE_VALUE': 255, 'NUM_CLASSES': 54, 'CONVS_DIM': 128, 'COMMON_STRIDE': 4, 'NORM': 'GN', 'LOSS_WEIGHT': 1.0}), 'PANOPTIC_FPN': CfgNode({'INSTANCE_LOSS_WEIGHT': 1.0, 'COMBINE': CfgNode({'ENABLED': True, 'OVERLAP_THRESH': 0.5, 'STUFF_AREA_LIMIT': 4096, 'INSTANCES_CONFIDENCE_THRESH': 0.5})}), 'RETINANET': CfgNode({'NUM_CLASSES': 80, 'IN_FEATURES': ['p3', 'p4', 'p5', 'p6', 'p7'], 'NUM_CONVS': 4, 'IOU_THRESHOLDS': [0.4, 0.5], 'IOU_LABELS': [0, -1, 1], 'PRIOR_PROB': 0.01, 'SCORE_THRESH_TEST': 0.05, 'TOPK_CANDIDATES_TEST': 1000, 'NMS_THRESH_TEST': 0.5, 'BBOX_REG_WEIGHTS': (1.0, 1.0, 1.0, 1.0), 'FOCAL_LOSS_GAMMA': 2.0, 'FOCAL_LOSS_ALPHA': 0.25, 'SMOOTH_L1_LOSS_BETA': 0.1, 'BBOX_REG_LOSS_TYPE': 'smooth_l1', 'NORM': ''}), 'RESNETS': CfgNode({'DEPTH': 50, 'OUT_FEATURES': ['res4'], 'NUM_GROUPS': 1, 'NORM': 'FrozenBN', 'WIDTH_PER_GROUP': 64, 'STRIDE_IN_1X1': True, 'RES5_DILATION': 1, 'RES2_OUT_CHANNELS': 256, 'STEM_OUT_CHANNELS': 64, 'DEFORM_ON_PER_STAGE': [False, False, False, False], 'DEFORM_MODULATED': False, 'DEFORM_NUM_GROUPS': 1}), 'FBNET_V2': CfgNode({'ARCH': 'FBNetV3_A', 'ARCH_DEF': [], 'STEM_IN_CHANNELS': 3, 'SCALE_FACTOR': 1.0, 'WIDTH_DIVISOR': 8, 'NORM': 'naiveSyncBN', 'NORM_ARGS': []}), 'VT_FPN': CfgNode({'IN_FEATURES': ['res2', 'res3', 'res4', 'res5'], 'OUT_CHANNELS': 256, 'LAYERS': 3, 'TOKEN_LS': [16, 16, 8, 8], 'TOKEN_C': 1024, 'HEADS': 16, 'MIN_GROUP_PLANES': 64, 'NORM': 'BN', 'POS_HWS': [], 'POS_N_DOWNSAMPLE': []}), 'FROZEN_LAYER_REG_EXP': [], 'KMEANS_ANCHORS': CfgNode({'KMEANS_ANCHORS_ON': False, 'NUM_CLUSTERS': 0, 'NUM_TRAINING_IMG': 0, 'DATASETS': (), 'RNG_SEED': 3}), 'SUBCLASS': CfgNode({'SUBCLASS_ON': False, 'NUM_SUBCLASSES': 0}), 'DDP_FIND_UNUSED_PARAMETERS': False}), 'INPUT': CfgNode({'MIN_SIZE_TRAIN': (224,), 'MIN_SIZE_TRAIN_SAMPLING': 'choice', 'MAX_SIZE_TRAIN': 320, 'MIN_SIZE_TEST': 224, 'MAX_SIZE_TEST': 320, 'RANDOM_FLIP': 'horizontal', 'CROP': CfgNode({'ENABLED': False, 'TYPE': 'relative_range', 'SIZE': [0.9, 0.9]}), 'FORMAT': 'BGR', 'MASK_FORMAT': 'polygon'}), 'DATASETS': CfgNode({'TRAIN': ('balloon_train',), 'PROPOSAL_FILES_TRAIN': (), 'PRECOMPUTED_PROPOSAL_TOPK_TRAIN': 2000, 'TEST': ('balloon_val',), 'PROPOSAL_FILES_TEST': (), 'PRECOMPUTED_PROPOSAL_TOPK_TEST': 1000, 'TRAIN_REPEAT_FACTOR': []}), 'DATALOADER': CfgNode({'NUM_WORKERS': 2, 'ASPECT_RATIO_GROUPING': True, 'SAMPLER_TRAIN': 'TrainingSampler', 'REPEAT_THRESHOLD': 0.0, 'FILTER_EMPTY_ANNOTATIONS': True}), 'SOLVER': CfgNode({'LR_SCHEDULER_NAME': 'WarmupCosineLR', 'MAX_ITER': 600, 'BASE_LR': 0.00025, 'MOMENTUM': 0.9, 'NESTEROV': False, 'WEIGHT_DECAY': 0.0001, 'WEIGHT_DECAY_NORM': 0.0, 'GAMMA': 0.1, 'STEPS': [], 'WARMUP_FACTOR': 0.001, 'WARMUP_ITERS': 1000, 'WARMUP_METHOD': 'linear', 'CHECKPOINT_PERIOD': 5000, 'IMS_PER_BATCH': 2, 'REFERENCE_WORLD_SIZE': 8, 'BIAS_LR_FACTOR': 1.0, 'WEIGHT_DECAY_BIAS': 0.0001, 'CLIP_GRADIENTS': CfgNode({'ENABLED': False, 'CLIP_TYPE': 'value', 'CLIP_VALUE': 1.0, 'NORM_TYPE': 2.0}), 'AMP': CfgNode({'ENABLED': False}), 'AUTO_SCALING_METHODS': ['default_scale_d2_configs', 'default_scale_quantization_configs'], 'OPTIMIZER': 'sgd', 'LR_MULTIPLIER_OVERWRITE': []}), 'TEST': CfgNode({'EXPECTED_RESULTS': [], 'EVAL_PERIOD': 10000, 'KEYPOINT_OKS_SIGMAS': [], 'DETECTIONS_PER_IMAGE': 100, 'AUG': CfgNode({'ENABLED': False, 'MIN_SIZES': (400, 500, 600, 700, 800, 900, 1000, 1100, 1200), 'MAX_SIZE': 4000, 'FLIP': True}), 'PRECISE_BN': CfgNode({'ENABLED': False, 'NUM_ITER': 200})}), 'OUTPUT_DIR': './output', 'SEED': -1, 'CUDNN_BENCHMARK': False, 'VIS_PERIOD': 0, 'GLOBAL': CfgNode({'HACK': 1.0}), 'MODEL_EMA': CfgNode({'ENABLED': False, 'DECAY': 0.9998, 'DEVICE': '', 'USE_EMA_WEIGHTS_FOR_EVAL_ONLY': False}), 'D2GO_DATA': CfgNode({'DATASETS': CfgNode({'TRAIN_CATEGORIES': (), 'TEST_CATEGORIES': (), 'COCO_INJECTION': CfgNode({'NAMES': [], 'IM_DIRS': [], 'JSON_FILES': []}), 'DYNAMIC_DATASETS': []}), 'MAPPER': CfgNode({'NAME': 'D2GoDatasetMapper', 'BACKFILL_SIZE': False, 'RETRY': 3, 'CATCH_EXCEPTION': True}), 'AUG_OPS': CfgNode({'TRAIN': ['ResizeShortestEdgeOp', 'RandomFlipOp'], 'TEST': ['ResizeShortestEdgeOp']}), 'TEST': CfgNode({'MAX_IMAGES': 0, 'SUBSET_SAMPLING': 'frontmost'})}), 'TENSORBOARD': CfgNode({'TRAIN_LOADER_VIS_WRITE_PERIOD': 20, 'TRAIN_LOADER_VIS_MAX_IMAGES': 16, 'TEST_VIS_MAX_IMAGES': 16}), 'QUANTIZATION': CfgNode({'EAGER_MODE': True, 'BACKEND': 'fbgemm', 'CUSTOM_QSCHEME': '', 'QAT': CfgNode({'ENABLED': False, 'BATCH_SIZE_FACTOR': 1.0, 'START_ITER': 35000, 'ENABLE_OBSERVER_ITER': 35000, 'DISABLE_OBSERVER_ITER': 38000, 'FREEZE_BN_ITER': 37000, 'UPDATE_OBSERVER_STATS_PERIODICALLY': False, 'UPDATE_OBSERVER_STATS_PERIOD': 1}), 'PTQ': CfgNode({'CALIBRATION_NUM_IMAGES': 1, 'CALIBRATION_FORCE_ON_GPU': False}), 'SILICON_QAT': CfgNode({'ENABLED': False})}), 'ABNORMAL_CHECKER': CfgNode({'ENABLED': False})})}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from d2go.runner import GeneralizedRCNNRunner\n",
    "\n",
    "\n",
    "def prepare_for_launch():\n",
    "    runner = GeneralizedRCNNRunner()\n",
    "    cfg = runner.get_default_cfg()\n",
    "    cfg.merge_from_file(model_zoo.get_config_file(\"faster_rcnn_fbnetv3a_C4.yaml\"))\n",
    "    cfg.MODEL_EMA.ENABLED = False\n",
    "    cfg.DATASETS.TRAIN = (\"balloon_train\",)\n",
    "    cfg.DATASETS.TEST = (\"balloon_val\",)\n",
    "    cfg.DATALOADER.NUM_WORKERS = 2\n",
    "    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(\"faster_rcnn_fbnetv3a_C4.yaml\")  # Let training initialize from model zoo\n",
    "    cfg.SOLVER.IMS_PER_BATCH = 2\n",
    "    cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR\n",
    "    cfg.SOLVER.MAX_ITER = 600    # 600 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset\n",
    "    cfg.SOLVER.STEPS = []        # do not decay learning rate\n",
    "    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, and good enough for this toy dataset (default: 512)\n",
    "    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # only has one class (ballon). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)\n",
    "    # NOTE: this config means the number of classes, but a few popular unofficial tutorials incorrect uses num_classes+1 here.\n",
    "    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
    "    return cfg, runner\n",
    "\n",
    "cfg, runner = prepare_for_launch()\n",
    "model = runner.build_model(cfg)\n",
    "runner.do_train(cfg, model, resume=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference & evaluation using the trained model\n",
    "Now, let's run inference with the trained model on the balloon validation dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading and preparing results...\n",
      "DONE (t=0.00s)\n",
      "creating index...\n",
      "index created!\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.494\n",
      " Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.651\n",
      " Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.543\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.104\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.757\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.204\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.526\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.526\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.118\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.810\n"
     ]
    }
   ],
   "source": [
    "metrics = runner.do_test(cfg, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The evaluation results are"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OrderedDict([('default', OrderedDict([('balloon_val', OrderedDict([('bbox', {'AP': 49.406689278871866, 'AP50': 65.11971524476927, 'AP75': 54.327369011410944, 'APs': 0.0, 'APm': 10.396039603960396, 'APl': 75.65561502020056})]))]))])\n"
     ]
    }
   ],
   "source": [
    "print(metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export to Int8 Model\n",
    "This section export int8 models using post-training quantization. For quantization-aware training, please see the [instructions](https://github.com/facebookresearch/d2go/tree/master/demo#quantization-aware-training)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[5m\u001b[31mWARNING\u001b[0m \u001b[32m[03/03 22:14:31 mobile_cv.arch.utils.helper]: \u001b[0mArguments ['width_divisor', 'dw_skip_bnrelu', 'zero_last_bn_gamma'] skipped for op Conv2d\n",
      "loading annotations into memory...\n",
      "Done (t=0.00s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/anaconda3/lib/python3.8/site-packages/torch/quantization/observer.py:121: UserWarning: Please use quant_min and quant_max to specify the range for observers.                     reduce_range will be deprecated in a future release of PyTorch.\n",
      "  warnings.warn(\n",
      "/home/ubuntu/anaconda3/lib/python3.8/site-packages/torch/tensor.py:565: UserWarning: floor_divide is deprecated, and will be removed in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values.\n",
      "To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor'). (Triggered internally at  /opt/conda/conda-bld/pytorch_1614499960647/work/aten/src/ATen/native/BinaryOps.cpp:341.)\n",
      "  return torch.floor_divide(self, other)\n",
      "/home/ubuntu/anaconda3/lib/python3.8/site-packages/torch/jit/_recursive.py:227: UserWarning: 'mask_on' was found in ScriptModule constants, but was not actually set in __init__. Consider removing it.\n",
      "  warnings.warn(\"'{}' was found in ScriptModule constants, \"\n",
      "/home/ubuntu/anaconda3/lib/python3.8/site-packages/torch/jit/_recursive.py:227: UserWarning: 'keypoint_on' was found in ScriptModule constants, but was not actually set in __init__. Consider removing it.\n",
      "  warnings.warn(\"'{}' was found in ScriptModule constants, \"\n",
      "/home/ubuntu/git/detectron2/detectron2/structures/image_list.py:92: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  assert t.shape[:-2] == tensors[0].shape[:-2], t.shape\n",
      "/home/ubuntu/git/detectron2/detectron2/structures/boxes.py:161: TracerWarning: torch.as_tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.\n",
      "  tensor = torch.as_tensor(tensor, dtype=torch.float32, device=device)\n",
      "/home/ubuntu/git/detectron2/detectron2/structures/boxes.py:166: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  assert tensor.dim() == 2 and tensor.size(-1) == 4, tensor.size()\n",
      "/home/ubuntu/git/detectron2/detectron2/modeling/proposal_generator/proposal_utils.py:101: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if not valid_mask.all():\n",
      "/home/ubuntu/git/detectron2/detectron2/structures/boxes.py:203: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  assert torch.isfinite(self.tensor).all(), \"Box tensor contains infinite or NaN!\"\n",
      "/home/ubuntu/anaconda3/lib/python3.8/site-packages/torch/tensor.py:593: RuntimeWarning: Iterating over a tensor might cause the trace to be incorrect. Passing a tensor of different shape won't change the number of iterations executed (and might lead to errors or silently give incorrect results).\n",
      "  warnings.warn('Iterating over a tensor might cause the trace to be incorrect. '\n",
      "/home/ubuntu/git/detectron2/detectron2/layers/nms.py:25: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  assert boxes.shape[-1] == 4\n",
      "/home/ubuntu/git/detectron2/detectron2/modeling/poolers.py:216: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  assert len(box_lists) == x[0].size(\n",
      "/home/ubuntu/git/detectron2/detectron2/layers/roi_align.py:55: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  assert rois.dim() == 2 and rois.size(1) == 5\n",
      "/home/ubuntu/git/d2go/d2go/modeling/backbone/fbnet_v2.py:372: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if len(x.shape) == 4 and (x.shape[2] > 1 or x.shape[3] > 1):\n",
      "/home/ubuntu/git/detectron2/detectron2/modeling/roi_heads/fast_rcnn.py:138: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if not valid_mask.all():\n",
      "/home/ubuntu/git/detectron2/detectron2/modeling/roi_heads/fast_rcnn.py:155: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if num_bbox_reg_classes == 1:\n"
     ]
    }
   ],
   "source": [
    "import copy\n",
    "from detectron2.data import build_detection_test_loader\n",
    "from d2go.export.api import convert_and_export_predictor\n",
    "from d2go.utils.testing.data_loader_helper import create_fake_detection_data_loader\n",
    "from d2go.export.d2_meta_arch import patch_d2_meta_arch\n",
    "\n",
    "import logging\n",
    "\n",
    "# disable all the warnings\n",
    "previous_level = logging.root.manager.disable\n",
    "logging.disable(logging.INFO)\n",
    "\n",
    "patch_d2_meta_arch()\n",
    "\n",
    "cfg_name = 'faster_rcnn_fbnetv3a_dsmask_C4.yaml'\n",
    "pytorch_model = model_zoo.get(cfg_name, trained=True, device='cpu')\n",
    "pytorch_model.eval()\n",
    "\n",
    "with create_fake_detection_data_loader(224, 320, is_train=False) as data_loader:\n",
    "    predictor_path = convert_and_export_predictor(\n",
    "            model_zoo.get_config(cfg_name),\n",
    "            copy.deepcopy(pytorch_model),\n",
    "            \"torchscript_int8\",\n",
    "            './',\n",
    "            data_loader,\n",
    "        )\n",
    "\n",
    "# recover the logging level\n",
    "logging.disable(previous_level)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the predictor using the exported int8 model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m[03/03 22:14:42 mobile_cv.predictor.api]: \u001b[0mCreating predictor from structured folder: ./torchscript_int8@tracing\n",
      "\u001b[32m[03/03 22:14:42 mobile_cv.predictor.api]: \u001b[0mLoading predictor info from ./torchscript_int8@tracing/predictor_info.json\n",
      "\u001b[32m[03/03 22:14:42 mobile_cv.predictor.model_wrappers]: \u001b[0mLoading /home/ubuntu/git/d2go/demo/torchscript_int8@tracing model from torchscript_int8@tracing ...\n"
     ]
    }
   ],
   "source": [
    "from mobile_cv.predictor.api import create_predictor\n",
    "model = create_predictor(predictor_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make predictions and Visualize the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f3b7404f430>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 216712 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "from d2go.utils.demo_predictor import DemoPredictor\n",
    "predictor = DemoPredictor(model)\n",
    "outputs = predictor(im)\n",
    "\n",
    "v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(\"coco_2017_train\"))\n",
    "out = v.draw_instance_predictions(outputs[\"instances\"].to(\"cpu\"))\n",
    "plt.imshow(out.get_image()[:, :, ::-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
