From d9eef066a2d55a45a159f1df7230e09c29bfcacb Mon Sep 17 00:00:00 2001
From: Lionel Weicker <lionel.weicker@arhs-spikeseed.com>
Date: Thu, 11 Apr 2024 13:50:43 +0200
Subject: [PATCH] Add reparameterization for yolov7-tiny

---
 tools/reparameterization.ipynb | 58 ++++++++++++++++++++++++++++++++--
 1 file changed, 56 insertions(+), 2 deletions(-)

diff --git a/tools/reparameterization.ipynb b/tools/reparameterization.ipynb
index 796b87c82d..e2a6ecd476 100644
--- a/tools/reparameterization.ipynb
+++ b/tools/reparameterization.ipynb
@@ -500,13 +500,67 @@
     "torch.save(ckpt, 'cfg/deploy/yolov7-e6e.pt')\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "1a1cbfa1",
+   "metadata": {},
+   "source": [
+    "## YOLOv7-tiny reparameterization"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "63a62625",
+   "id": "aebfde10",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "# import\n",
+    "from copy import deepcopy\n",
+    "from models.yolo import Model\n",
+    "import torch\n",
+    "from utils.torch_utils import select_device, is_parallel\n",
+    "import yaml\n",
+    "\n",
+    "device = select_device('0', batch_size=1)\n",
+    "# model trained by cfg/training/*.yaml\n",
+    "ckpt = torch.load('cfg/training/yolov7-tiny_training.pt', map_location=device)\n",
+    "# reparameterized model in cfg/deploy/*.yaml\n",
+    "model = Model('cfg/deploy/yolov7-tiny.yaml', ch=3, nc=80).to(device)\n",
+    "\n",
+    "with open('cfg/deploy/yolov7-tiny.yaml') as f:\n",
+    "    yml = yaml.load(f, Loader=yaml.SafeLoader)\n",
+    "anchors = len(yml['anchors'][0]) // 2\n",
+    "\n",
+    "# copy intersect weights\n",
+    "state_dict = ckpt['model'].float().state_dict()\n",
+    "exclude = []\n",
+    "intersect_state_dict = {k: v for k, v in state_dict.items() if k in model.state_dict() and not any(x in k for x in exclude) and v.shape == model.state_dict()[k].shape}\n",
+    "model.load_state_dict(intersect_state_dict, strict=False)\n",
+    "model.names = ckpt['model'].names\n",
+    "model.nc = ckpt['model'].nc\n",
+    "\n",
+    "# reparametrized YOLOR\n",
+    "for i in range((model.nc+5)*anchors):\n",
+    "    model.state_dict()['model.77.m.0.weight'].data[i, :, :, :] *= state_dict['model.77.im.0.implicit'].data[:, i, : :].squeeze()\n",
+    "    model.state_dict()['model.77.m.1.weight'].data[i, :, :, :] *= state_dict['model.77.im.1.implicit'].data[:, i, : :].squeeze()\n",
+    "    model.state_dict()['model.77.m.2.weight'].data[i, :, :, :] *= state_dict['model.77.im.2.implicit'].data[:, i, : :].squeeze()\n",
+    "model.state_dict()['model.77.m.0.bias'].data += state_dict['model.77.m.0.weight'].mul(state_dict['model.77.ia.0.implicit']).sum(1).squeeze()\n",
+    "model.state_dict()['model.77.m.1.bias'].data += state_dict['model.77.m.1.weight'].mul(state_dict['model.77.ia.1.implicit']).sum(1).squeeze()\n",
+    "model.state_dict()['model.77.m.2.bias'].data += state_dict['model.77.m.2.weight'].mul(state_dict['model.77.ia.2.implicit']).sum(1).squeeze()\n",
+    "model.state_dict()['model.77.m.0.bias'].data *= state_dict['model.77.im.0.implicit'].data.squeeze()\n",
+    "model.state_dict()['model.77.m.1.bias'].data *= state_dict['model.77.im.1.implicit'].data.squeeze()\n",
+    "model.state_dict()['model.77.m.2.bias'].data *= state_dict['model.77.im.2.implicit'].data.squeeze()\n",
+    "\n",
+    "# model to be saved\n",
+    "ckpt = {'model': deepcopy(model.module if is_parallel(model) else model).half(),\n",
+    "        'optimizer': None,\n",
+    "        'training_results': None,\n",
+    "        'epoch': -1}\n",
+    "\n",
+    "# save reparameterized model\n",
+    "torch.save(ckpt, 'cfg/deploy/yolov7-tiny.pt')"
+   ]
   }
  ],
  "metadata": {