1.yolo2_pytorch_onnx_save_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# pytorch-yolo2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### create the subfolder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir onnx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ref: https://github.com/longcw/yolo2-pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### download weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! wget http://pjreddie.com/media/files/yolo.weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### try predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "layer     filters    size              input                output\n",
      "    0 conv     32  3 x 3 / 1   416 x 416 x   3   ->   416 x 416 x  32\n",
      "    1 max          2 x 2 / 2   416 x 416 x  32   ->   208 x 208 x  32\n",
      "    2 conv     64  3 x 3 / 1   208 x 208 x  32   ->   208 x 208 x  64\n",
      "    3 max          2 x 2 / 2   208 x 208 x  64   ->   104 x 104 x  64\n",
      "    4 conv    128  3 x 3 / 1   104 x 104 x  64   ->   104 x 104 x 128\n",
      "    5 conv     64  1 x 1 / 1   104 x 104 x 128   ->   104 x 104 x  64\n",
      "    6 conv    128  3 x 3 / 1   104 x 104 x  64   ->   104 x 104 x 128\n",
      "    7 max          2 x 2 / 2   104 x 104 x 128   ->    52 x  52 x 128\n",
      "    8 conv    256  3 x 3 / 1    52 x  52 x 128   ->    52 x  52 x 256\n",
      "    9 conv    128  1 x 1 / 1    52 x  52 x 256   ->    52 x  52 x 128\n",
      "   10 conv    256  3 x 3 / 1    52 x  52 x 128   ->    52 x  52 x 256\n",
      "   11 max          2 x 2 / 2    52 x  52 x 256   ->    26 x  26 x 256\n",
      "   12 conv    512  3 x 3 / 1    26 x  26 x 256   ->    26 x  26 x 512\n",
      "   13 conv    256  1 x 1 / 1    26 x  26 x 512   ->    26 x  26 x 256\n",
      "   14 conv    512  3 x 3 / 1    26 x  26 x 256   ->    26 x  26 x 512\n",
      "   15 conv    256  1 x 1 / 1    26 x  26 x 512   ->    26 x  26 x 256\n",
      "   16 conv    512  3 x 3 / 1    26 x  26 x 256   ->    26 x  26 x 512\n",
      "   17 max          2 x 2 / 2    26 x  26 x 512   ->    13 x  13 x 512\n",
      "   18 conv   1024  3 x 3 / 1    13 x  13 x 512   ->    13 x  13 x1024\n",
      "   19 conv    512  1 x 1 / 1    13 x  13 x1024   ->    13 x  13 x 512\n",
      "   20 conv   1024  3 x 3 / 1    13 x  13 x 512   ->    13 x  13 x1024\n",
      "   21 conv    512  1 x 1 / 1    13 x  13 x1024   ->    13 x  13 x 512\n",
      "   22 conv   1024  3 x 3 / 1    13 x  13 x 512   ->    13 x  13 x1024\n",
      "   23 conv   1024  3 x 3 / 1    13 x  13 x1024   ->    13 x  13 x1024\n",
      "   24 conv   1024  3 x 3 / 1    13 x  13 x1024   ->    13 x  13 x1024\n",
      "   25 route  16\n",
      "   26 conv     64  1 x 1 / 1    26 x  26 x 512   ->    26 x  26 x  64\n",
      "   27 reorg              / 2    26 x  26 x  64   ->    13 x  13 x 256\n",
      "   28 route  27 24\n",
      "   29 conv   1024  3 x 3 / 1    13 x  13 x1280   ->    13 x  13 x1024\n",
      "   30 conv    425  1 x 1 / 1    13 x  13 x1024   ->    13 x  13 x 425\n",
      "   31 detection\n",
      "/home/phoebehuang/Desktop/yolo2_onnx/cfg.py:175: UserWarning: src is not broadcastable to dst, but they have the same number of elements.  Falling back to deprecated pointwise behavior.\n",
      "  conv_model.weight.data.copy_(torch.from_numpy(buf[start:start+num_w])); start = start + num_w\n",
      "/home/phoebehuang/Desktop/yolo2_onnx/cfg.py:157: UserWarning: src is not broadcastable to dst, but they have the same number of elements.  Falling back to deprecated pointwise behavior.\n",
      "  conv_model.weight.data.copy_(torch.from_numpy(buf[start:start+num_w])); start = start + num_w\n",
      "Loading weights from ./yolo.weights... Done!\n",
      "/home/phoebehuang/Desktop/yolo2_onnx/utils.py:140: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "  cls_confs = torch.nn.Softmax()(Variable(output[5:5+num_classes].transpose(0,1))).data\n",
      "./data/dog.jpg: Predicted in 0.011061 seconds.\n",
      "truck: 0.934710\n",
      "bicycle: 0.998012\n",
      "dog: 0.990524\n",
      "save plot results to predictions.jpg\n"
     ]
    }
   ],
   "source": [
    "! python detect.py './cfg/yolo.cfg' './yolo.weights' './data/dog.jpg'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### get model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from darknet import Darknet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfgfile =  './cfg/yolo.cfg' \n",
    "weightfile =  './yolo.weights'\n",
    "imgfile = './data/dog.jpg'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading weights from ./yolo.weights... Done!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "cfg.py:175: UserWarning: src is not broadcastable to dst, but they have the same number of elements.  Falling back to deprecated pointwise behavior.\n",
      "  conv_model.weight.data.copy_(torch.from_numpy(buf[start:start+num_w])); start = start + num_w\n",
      "cfg.py:157: UserWarning: src is not broadcastable to dst, but they have the same number of elements.  Falling back to deprecated pointwise behavior.\n",
      "  conv_model.weight.data.copy_(torch.from_numpy(buf[start:start+num_w])); start = start + num_w\n"
     ]
    }
   ],
   "source": [
    "m = Darknet(cfgfile)\n",
    "m.load_weights(weightfile)\n",
    "print('Loading weights from %s... Done!' % (weightfile))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### save detection information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "op_dict = {\n",
    "    'num_classes':m.num_classes,\n",
    "    'anchors':m.anchors,\n",
    "    'num_anchors':m.num_anchors\n",
    "}\n",
    "pickle.dump(op_dict, open('detection_information.pkl','wb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# use Onnx to convert model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ref: https://github.com/onnx/tutorials/blob/master/tutorials/PytorchOnnxExport.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.onnx\n",
    "from torch.autograd import Variable\n",
    "\n",
    "# Standard ImageNet input - 3 channels, 224x224,\n",
    "# values don't matter as we care about network structure.\n",
    "# But they can also be real inputs.\n",
    "dummy_input = Variable(torch.randn(1, 3, 416, 416))\n",
    "# Obtain your model, it can be also constructed in your script explicitly\n",
    "model = m\n",
    "# Invoke export\n",
    "torch.onnx.export(model, dummy_input, \"onnx/yolo2.onnx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}