mindspore_yolo/tools/reparameterization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d7cbe5ee",
   "metadata": {},
   "source": [
    "# Reparameterization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13393b70",
   "metadata": {},
   "source": [
    "## YOLOv7 reparameterization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf53becf",
   "metadata": {},
   "outputs": [],
   "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",
    "\n",
    "device = select_device('0', batch_size=1)\n",
    "# model trained by cfg/training/*.yaml\n",
    "ckpt = torch.load('cfg/training/yolov7.pt', map_location=device)\n",
    "# reparameterized model in cfg/deploy/*.yaml\n",
    "model = Model('cfg/deploy/yolov7.yaml', ch=3, nc=80).to(device)\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(255):\n",
    "    model.state_dict()['model.105.m.0.weight'].data[i, :, :, :] *= state_dict['model.105.im.0.implicit'].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.105.m.1.weight'].data[i, :, :, :] *= state_dict['model.105.im.1.implicit'].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.105.m.2.weight'].data[i, :, :, :] *= state_dict['model.105.im.2.implicit'].data[:, i, : :].squeeze()\n",
    "model.state_dict()['model.105.m.0.bias'].data += state_dict['model.105.m.0.weight'].mul(state_dict['model.105.ia.0.implicit']).sum(1).squeeze()\n",
    "model.state_dict()['model.105.m.1.bias'].data += state_dict['model.105.m.1.weight'].mul(state_dict['model.105.ia.1.implicit']).sum(1).squeeze()\n",
    "model.state_dict()['model.105.m.2.bias'].data += state_dict['model.105.m.2.weight'].mul(state_dict['model.105.ia.2.implicit']).sum(1).squeeze()\n",
    "model.state_dict()['model.105.m.0.bias'].data *= state_dict['model.105.im.0.implicit'].data.squeeze()\n",
    "model.state_dict()['model.105.m.1.bias'].data *= state_dict['model.105.im.1.implicit'].data.squeeze()\n",
    "model.state_dict()['model.105.m.2.bias'].data *= state_dict['model.105.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.pt')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b396a53",
   "metadata": {},
   "source": [
    "## YOLOv7x reparameterization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d54d17f",
   "metadata": {},
   "outputs": [],
   "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",
    "\n",
    "device = select_device('0', batch_size=1)\n",
    "# model trained by cfg/training/*.yaml\n",
    "ckpt = torch.load('cfg/training/yolov7x.pt', map_location=device)\n",
    "# reparameterized model in cfg/deploy/*.yaml\n",
    "model = Model('cfg/deploy/yolov7x.yaml', ch=3, nc=80).to(device)\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(255):\n",
    "    model.state_dict()['model.121.m.0.weight'].data[i, :, :, :] *= state_dict['model.121.im.0.implicit'].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.121.m.1.weight'].data[i, :, :, :] *= state_dict['model.121.im.1.implicit'].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.121.m.2.weight'].data[i, :, :, :] *= state_dict['model.121.im.2.implicit'].data[:, i, : :].squeeze()\n",
    "model.state_dict()['model.121.m.0.bias'].data += state_dict['model.121.m.0.weight'].mul(state_dict['model.121.ia.0.implicit']).sum(1).squeeze()\n",
    "model.state_dict()['model.121.m.1.bias'].data += state_dict['model.121.m.1.weight'].mul(state_dict['model.121.ia.1.implicit']).sum(1).squeeze()\n",
    "model.state_dict()['model.121.m.2.bias'].data += state_dict['model.121.m.2.weight'].mul(state_dict['model.121.ia.2.implicit']).sum(1).squeeze()\n",
    "model.state_dict()['model.121.m.0.bias'].data *= state_dict['model.121.im.0.implicit'].data.squeeze()\n",
    "model.state_dict()['model.121.m.1.bias'].data *= state_dict['model.121.im.1.implicit'].data.squeeze()\n",
    "model.state_dict()['model.121.m.2.bias'].data *= state_dict['model.121.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/yolov7x.pt')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11a9108e",
   "metadata": {},
   "source": [
    "## YOLOv7-W6 reparameterization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d032c629",
   "metadata": {},
   "outputs": [],
   "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",
    "\n",
    "device = select_device('0', batch_size=1)\n",
    "# model trained by cfg/training/*.yaml\n",
    "ckpt = torch.load('cfg/training/yolov7-w6.pt', map_location=device)\n",
    "# reparameterized model in cfg/deploy/*.yaml\n",
    "model = Model('cfg/deploy/yolov7-w6.yaml', ch=3, nc=80).to(device)\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",
    "idx = 118\n",
    "idx2 = 122\n",
    "\n",
    "# copy weights of lead head\n",
    "model.state_dict()['model.{}.m.0.weight'.format(idx)].data -= model.state_dict()['model.{}.m.0.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.1.weight'.format(idx)].data -= model.state_dict()['model.{}.m.1.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.2.weight'.format(idx)].data -= model.state_dict()['model.{}.m.2.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.3.weight'.format(idx)].data -= model.state_dict()['model.{}.m.3.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.0.weight'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.1.weight'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.2.weight'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.3.weight'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data -= model.state_dict()['model.{}.m.0.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data -= model.state_dict()['model.{}.m.1.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data -= model.state_dict()['model.{}.m.2.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data -= model.state_dict()['model.{}.m.3.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.bias'.format(idx2)].data\n",
    "\n",
    "# reparametrized YOLOR\n",
    "for i in range(255):\n",
    "    model.state_dict()['model.{}.m.0.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.0.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.1.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.1.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.2.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.2.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.3.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.3.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].mul(state_dict['model.{}.ia.0.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].mul(state_dict['model.{}.ia.1.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].mul(state_dict['model.{}.ia.2.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].mul(state_dict['model.{}.ia.3.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data *= state_dict['model.{}.im.0.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data *= state_dict['model.{}.im.1.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data *= state_dict['model.{}.im.2.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data *= state_dict['model.{}.im.3.implicit'.format(idx2)].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-w6.pt')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f093d43",
   "metadata": {},
   "source": [
    "## YOLOv7-E6 reparameterization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa2b2142",
   "metadata": {},
   "outputs": [],
   "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",
    "\n",
    "device = select_device('0', batch_size=1)\n",
    "# model trained by cfg/training/*.yaml\n",
    "ckpt = torch.load('cfg/training/yolov7-e6.pt', map_location=device)\n",
    "# reparameterized model in cfg/deploy/*.yaml\n",
    "model = Model('cfg/deploy/yolov7-e6.yaml', ch=3, nc=80).to(device)\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",
    "idx = 140\n",
    "idx2 = 144\n",
    "\n",
    "# copy weights of lead head\n",
    "model.state_dict()['model.{}.m.0.weight'.format(idx)].data -= model.state_dict()['model.{}.m.0.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.1.weight'.format(idx)].data -= model.state_dict()['model.{}.m.1.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.2.weight'.format(idx)].data -= model.state_dict()['model.{}.m.2.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.3.weight'.format(idx)].data -= model.state_dict()['model.{}.m.3.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.0.weight'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.1.weight'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.2.weight'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.3.weight'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data -= model.state_dict()['model.{}.m.0.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data -= model.state_dict()['model.{}.m.1.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data -= model.state_dict()['model.{}.m.2.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data -= model.state_dict()['model.{}.m.3.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.bias'.format(idx2)].data\n",
    "\n",
    "# reparametrized YOLOR\n",
    "for i in range(255):\n",
    "    model.state_dict()['model.{}.m.0.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.0.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.1.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.1.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.2.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.2.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.3.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.3.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].mul(state_dict['model.{}.ia.0.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].mul(state_dict['model.{}.ia.1.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].mul(state_dict['model.{}.ia.2.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].mul(state_dict['model.{}.ia.3.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data *= state_dict['model.{}.im.0.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data *= state_dict['model.{}.im.1.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data *= state_dict['model.{}.im.2.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data *= state_dict['model.{}.im.3.implicit'.format(idx2)].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-e6.pt')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3bccf89",
   "metadata": {},
   "source": [
    "## YOLOv7-D6 reparameterization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5216b70",
   "metadata": {},
   "outputs": [],
   "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",
    "\n",
    "device = select_device('0', batch_size=1)\n",
    "# model trained by cfg/training/*.yaml\n",
    "ckpt = torch.load('cfg/training/yolov7-d6.pt', map_location=device)\n",
    "# reparameterized model in cfg/deploy/*.yaml\n",
    "model = Model('cfg/deploy/yolov7-d6.yaml', ch=3, nc=80).to(device)\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",
    "idx = 162\n",
    "idx2 = 166\n",
    "\n",
    "# copy weights of lead head\n",
    "model.state_dict()['model.{}.m.0.weight'.format(idx)].data -= model.state_dict()['model.{}.m.0.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.1.weight'.format(idx)].data -= model.state_dict()['model.{}.m.1.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.2.weight'.format(idx)].data -= model.state_dict()['model.{}.m.2.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.3.weight'.format(idx)].data -= model.state_dict()['model.{}.m.3.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.0.weight'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.1.weight'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.2.weight'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.3.weight'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data -= model.state_dict()['model.{}.m.0.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data -= model.state_dict()['model.{}.m.1.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data -= model.state_dict()['model.{}.m.2.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data -= model.state_dict()['model.{}.m.3.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.bias'.format(idx2)].data\n",
    "\n",
    "# reparametrized YOLOR\n",
    "for i in range(255):\n",
    "    model.state_dict()['model.{}.m.0.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.0.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.1.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.1.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.2.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.2.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.3.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.3.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].mul(state_dict['model.{}.ia.0.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].mul(state_dict['model.{}.ia.1.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].mul(state_dict['model.{}.ia.2.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].mul(state_dict['model.{}.ia.3.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data *= state_dict['model.{}.im.0.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data *= state_dict['model.{}.im.1.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data *= state_dict['model.{}.im.2.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data *= state_dict['model.{}.im.3.implicit'.format(idx2)].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-d6.pt')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "334c273b",
   "metadata": {},
   "source": [
    "## YOLOv7-E6E reparameterization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "635fd8d2",
   "metadata": {},
   "outputs": [],
   "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",
    "\n",
    "device = select_device('0', batch_size=1)\n",
    "# model trained by cfg/training/*.yaml\n",
    "ckpt = torch.load('cfg/training/yolov7-e6e.pt', map_location=device)\n",
    "# reparameterized model in cfg/deploy/*.yaml\n",
    "model = Model('cfg/deploy/yolov7-e6e.yaml', ch=3, nc=80).to(device)\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",
    "idx = 261\n",
    "idx2 = 265\n",
    "\n",
    "# copy weights of lead head\n",
    "model.state_dict()['model.{}.m.0.weight'.format(idx)].data -= model.state_dict()['model.{}.m.0.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.1.weight'.format(idx)].data -= model.state_dict()['model.{}.m.1.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.2.weight'.format(idx)].data -= model.state_dict()['model.{}.m.2.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.3.weight'.format(idx)].data -= model.state_dict()['model.{}.m.3.weight'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.0.weight'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.1.weight'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.2.weight'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.3.weight'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data -= model.state_dict()['model.{}.m.0.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data -= model.state_dict()['model.{}.m.1.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data -= model.state_dict()['model.{}.m.2.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data -= model.state_dict()['model.{}.m.3.bias'.format(idx)].data\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.bias'.format(idx2)].data\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.bias'.format(idx2)].data\n",
    "\n",
    "# reparametrized YOLOR\n",
    "for i in range(255):\n",
    "    model.state_dict()['model.{}.m.0.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.0.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.1.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.1.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.2.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.2.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "    model.state_dict()['model.{}.m.3.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.3.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].mul(state_dict['model.{}.ia.0.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].mul(state_dict['model.{}.ia.1.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].mul(state_dict['model.{}.ia.2.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].mul(state_dict['model.{}.ia.3.implicit'.format(idx2)]).sum(1).squeeze()\n",
    "model.state_dict()['model.{}.m.0.bias'.format(idx)].data *= state_dict['model.{}.im.0.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.1.bias'.format(idx)].data *= state_dict['model.{}.im.1.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.2.bias'.format(idx)].data *= state_dict['model.{}.im.2.implicit'.format(idx2)].data.squeeze()\n",
    "model.state_dict()['model.{}.m.3.bias'.format(idx)].data *= state_dict['model.{}.im.3.implicit'.format(idx2)].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-e6e.pt')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63a62625",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}