Upload 2 files

QLNet_symmetry.ipynb

@@ -0,0 +1,551 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "71b6152c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch, timm\n",
+    "from qlnet import QLNet"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "4e7ed219",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "m = QLNet()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "3f703be8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "state_dict = torch.load('qlnet-50-v0.pth.tar')['state_dict']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "435e2358",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "m.load_state_dict(state_dict)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "f14d984a",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "QLNet(\n",
+       "  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
+       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "  (act1): ReLU(inplace=True)\n",
+       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
+       "  (layer1): Sequential(\n",
+       "    (0): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(64, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=256, bias=False)\n",
+       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(512, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (1): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(64, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=256, bias=False)\n",
+       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(512, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (2): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(64, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=256, bias=False)\n",
+       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(512, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "  )\n",
+       "  (layer2): Sequential(\n",
+       "    (0): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(64, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=256, bias=False)\n",
+       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): ConvBN(\n",
+       "        (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
+       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (1): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(128, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
+       "      (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(1024, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (2): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(128, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
+       "      (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(1024, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (3): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(128, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
+       "      (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(1024, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "  )\n",
+       "  (layer3): Sequential(\n",
+       "    (0): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(128, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(512, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=512, bias=False)\n",
+       "      (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): ConvBN(\n",
+       "        (conv): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
+       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (1): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
+       "      (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (2): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
+       "      (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (3): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
+       "      (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (4): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
+       "      (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (5): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
+       "      (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "  )\n",
+       "  (layer4): Sequential(\n",
+       "    (0): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(512, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=512, bias=False)\n",
+       "      (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): ConvBN(\n",
+       "        (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
+       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (1): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(1024, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024, bias=False)\n",
+       "      (bn2): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "    (2): QLBlock(\n",
+       "      (conv1): ConvBN(\n",
+       "        (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (conv2): Conv2d(1024, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024, bias=False)\n",
+       "      (bn2): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): ConvBN(\n",
+       "        (conv): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (skip): Identity()\n",
+       "      (act3): hardball()\n",
+       "    )\n",
+       "  )\n",
+       "  (act): hardball()\n",
+       "  (global_pool): SelectAdaptivePool2d (pool_type=avg, flatten=Flatten(start_dim=1, end_dim=-1))\n",
+       "  (fc): Linear(in_features=512, out_features=1000, bias=True)\n",
+       ")"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "m.eval()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "2099b937",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "layer1 >>\n",
+      "torch.Size([512, 64, 1, 1])\n",
+      "torch.Size([64, 512, 1, 1])\n",
+      "torch.Size([512, 64, 1, 1])\n",
+      "torch.Size([64, 512, 1, 1])\n",
+      "torch.Size([512, 64, 1, 1])\n",
+      "torch.Size([64, 512, 1, 1])\n",
+      "layer2 >>\n",
+      "torch.Size([512, 64, 1, 1])\n",
+      "torch.Size([128, 512, 1, 1])\n",
+      "torch.Size([128, 64, 1, 1])\n",
+      "torch.Size([1024, 128, 1, 1])\n",
+      "torch.Size([128, 1024, 1, 1])\n",
+      "torch.Size([1024, 128, 1, 1])\n",
+      "torch.Size([128, 1024, 1, 1])\n",
+      "torch.Size([1024, 128, 1, 1])\n",
+      "torch.Size([128, 1024, 1, 1])\n",
+      "layer3 >>\n",
+      "torch.Size([1024, 128, 1, 1])\n",
+      "torch.Size([256, 1024, 1, 1])\n",
+      "torch.Size([256, 128, 1, 1])\n",
+      "torch.Size([1024, 256, 1, 1])\n",
+      "torch.Size([256, 1024, 1, 1])\n",
+      "torch.Size([1024, 256, 1, 1])\n",
+      "torch.Size([256, 1024, 1, 1])\n",
+      "torch.Size([1024, 256, 1, 1])\n",
+      "torch.Size([256, 1024, 1, 1])\n",
+      "torch.Size([1024, 256, 1, 1])\n",
+      "torch.Size([256, 1024, 1, 1])\n",
+      "torch.Size([1024, 256, 1, 1])\n",
+      "torch.Size([256, 1024, 1, 1])\n",
+      "layer4 >>\n",
+      "torch.Size([1024, 256, 1, 1])\n",
+      "torch.Size([512, 1024, 1, 1])\n",
+      "torch.Size([512, 256, 1, 1])\n",
+      "torch.Size([2048, 512, 1, 1])\n",
+      "torch.Size([512, 2048, 1, 1])\n",
+      "torch.Size([2048, 512, 1, 1])\n",
+      "torch.Size([512, 2048, 1, 1])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# fuse ConvBN\n",
+    "i = 1\n",
+    "for layer in [m.layer1, m.layer2, m.layer3, m.layer4]:\n",
+    "    print(f'layer{i} >>')\n",
+    "    for block in layer:\n",
+    "        # Fuse the weights in conv1 and conv3\n",
+    "        block.conv1.fuse_bn()\n",
+    "        print(block.conv1.fused_weight.size())\n",
+    "        block.conv3.fuse_bn()\n",
+    "        print(block.conv3.fused_weight.size())\n",
+    "        if not isinstance(block.skip, torch.nn.Identity):\n",
+    "            layer[0].skip.fuse_bn()\n",
+    "            print(layer[0].skip.fused_weight.size())\n",
+    "    i += 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "b3a55f82",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inpt = torch.randn(5,3,224,224)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "dccbf19c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "out_old = m(inpt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "f0c74a04",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([5, 1000])"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "out_old.size()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "a5991c8f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def apply_transform(block1, block2, Q, keep_identity=True):\n",
+    "    with torch.no_grad():\n",
+    "        # Ensure that the out_channels of block1 is equal to the in_channels of block2\n",
+    "        assert Q.size()[0] == Q.size()[1], \"Q needs to be a square matrix\"\n",
+    "        n = Q.size()[0]\n",
+    "        assert block1.conv3.conv.out_channels == n and block2.conv1.conv.in_channels == n, \"Mismatched channels between blocks\"\n",
+    "\n",
+    "        n = block1.conv3.conv.out_channels\n",
+    "        \n",
+    "        # Calculate the inverse of Q\n",
+    "        Q_inv = torch.inverse(Q)\n",
+    "\n",
+    "        # Modify the weights of conv layers in block1\n",
+    "        block1.conv3.fused_weight.data = torch.einsum('ij,jklm->iklm', Q, block1.conv3.fused_weight.data)\n",
+    "        block1.conv3.fused_bias.data = torch.einsum('ij,j->i', Q, block1.conv3.fused_bias.data)\n",
+    "        \n",
+    "        if isinstance(block1.skip, torch.nn.Identity):\n",
+    "            if not keep_identity:\n",
+    "                block1.skip = torch.nn.Conv2d(n, n, kernel_size=1, bias=False)\n",
+    "                block1.skip.weight.data = Q.unsqueeze(-1).unsqueeze(-1)\n",
+    "        else:\n",
+    "            block1.skip.fused_weight.data = torch.einsum('ij,jklm->iklm', Q, block1.skip.fused_weight.data)\n",
+    "            block1.skip.fused_bias.data = torch.einsum('ij,j->i', Q, block1.skip.fused_bias.data)\n",
+    "\n",
+    "        # Modify the weights of conv layers in block2\n",
+    "        block2.conv1.fused_weight.data = torch.einsum('ki,jklm->jilm', Q_inv, block2.conv1.fused_weight.data)\n",
+    "        \n",
+    "        if isinstance(block2.skip, torch.nn.Identity):\n",
+    "            if not keep_identity:\n",
+    "                block2.skip = torch.nn.Conv2d(n, n, kernel_size=1, bias=False)\n",
+    "                block2.skip.weight.data = Q_inv.unsqueeze(-1).unsqueeze(-1)\n",
+    "        else:\n",
+    "            block2.skip.fused_weight.data = torch.einsum('ki,jklm->jilm', Q_inv, block2.skip.fused_weight.data)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "dd96acd7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "Q = torch.nn.init.orthogonal_(torch.empty(256, 256))\n",
+    "for i in range(5):\n",
+    "    apply_transform(m.layer3[i], m.layer3[i+1], Q, True)\n",
+    "apply_transform(m.layer3[5], m.layer4[0], Q, True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "e5d3628d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6.0558319091796875e-05\n"
+     ]
+    }
+   ],
+   "source": [
+    "out_new = m(inpt)\n",
+    "print((out_new - out_old).abs().max().item())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9fce3a38",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a54fe8b",
+   "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.10.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

qlnet.py

@@ -0,0 +1,385 @@
+"""PyTorch ResNet
+
+This started as a copy of https://github.com/pytorch/vision 'resnet.py' (BSD-3-Clause) with
+additional dropout and dynamic global avg/max pool.
+
+ResNeXt, SE-ResNeXt, SENet, and MXNet Gluon stem/downsample variants, tiered stems added by Ross Wightman
+
+Copyright 2019, Ross Wightman
+"""
+import math 
+from functools import partial
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.layers import DropBlock2d, DropPath, AvgPool2dSame, BlurPool2d, GroupNorm, create_attn, get_attn, \
+    get_act_layer, get_norm_layer, create_classifier, LayerNorm2d
+
+
+def get_padding(kernel_size, stride, dilation=1):
+    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
+    return padding
+
+
+class softball(nn.Module):
+    def __init__(self, radius2=None, inplace=True):
+        super(softball, self).__init__()
+        self.radius2 = radius2 if radius2 is not None else None
+
+    def forward(self, x):
+        if self.radius2 is None:
+            self.radius2 = x.size()[1]
+        norm = torch.sqrt(1 + (x*x).sum(1, keepdim=True) / self.radius2)
+        return x / norm
+
+class hardball(nn.Module):
+    def __init__(self, radius2=None):
+        super(hardball, self).__init__()
+        self.radius = np.sqrt(radius2) if radius2 is not None else None
+
+    def forward(self, x):
+        norm = torch.sqrt((x*x).sum(1, keepdim=True))
+        if self.radius is None:
+            self.radius = np.sqrt(x.size()[1])
+        return torch.where(norm > self.radius, self.radius * x / norm, x)
+
+
+class ConvBN(nn.Module):
+    def __init__(self, conv, bn):
+        super(ConvBN, self).__init__()
+        self.conv = conv
+        self.bn = bn
+        self.fused_weight = None
+        self.fused_bias = None
+
+    def forward(self, x):
+        if self.training:
+            x = self.conv(x)
+            x = self.bn(x)
+        else:
+            if self.fused_weight is not None and self.fused_bias is not None:
+                x = F.conv2d(x, self.fused_weight, self.fused_bias, 
+                            self.conv.stride, self.conv.padding, 
+                            self.conv.dilation, self.conv.groups)
+            else:
+                x = self.conv(x)
+                x = self.bn(x)
+        return x
+
+    def fuse_bn(self):
+        if self.training:
+            raise RuntimeError("Call fuse_bn only in eval mode")
+        
+        # Calculate the fused weight and bias
+        w = self.conv.weight
+        mean = self.bn.running_mean
+        var = torch.sqrt(self.bn.running_var + self.bn.eps)
+        gamma = self.bn.weight
+        beta = self.bn.bias
+
+        self.fused_weight = w * (gamma / var).reshape(-1, 1, 1, 1)
+        self.fused_bias = beta - (gamma * mean / var)
+
+
+class QLBlock(nn.Module): # quasilinear hyperbolic system
+    expansion = 1
+
+    def __init__(
+            self,
+            inplanes,
+            planes,
+            stride=1,
+            downsample=None,
+            cardinality=1,
+            base_width=64,
+            reduce_first=1,
+            dilation=1,
+            first_dilation=None,
+            act_layer=nn.ReLU,
+            norm_layer=nn.BatchNorm2d,
+    ):
+        super(QLBlock, self).__init__()
+
+        k = 4 if inplanes <= 128 else 2
+        width = inplanes * k
+        outplanes = inplanes if downsample is None else inplanes * 2
+        first_dilation = first_dilation or dilation
+
+        self.conv1 = ConvBN(
+            nn.Conv2d(inplanes, width*2, kernel_size=1, stride=1,
+                dilation=first_dilation, groups=1, bias=False),
+            norm_layer(width*2))
+
+        self.conv2 = nn.Conv2d(width, width*2, kernel_size=3, stride=stride,
+                padding=1, dilation=first_dilation, groups=width, bias=False)
+        self.bn2 = norm_layer(width*2)
+
+        self.conv3 = ConvBN(
+            nn.Conv2d(width*2, outplanes, kernel_size=1, groups=1, bias=False),
+            norm_layer(outplanes))
+
+        self.skip = ConvBN(
+            nn.Conv2d(inplanes, outplanes, kernel_size=1, stride=stride,
+                dilation=first_dilation, groups=1, bias=False),
+            norm_layer(outplanes)) if downsample is not None else nn.Identity()
+
+        self.act3 = hardball(radius2=outplanes) # if downsample is not None else None
+
+    def zero_init_last(self):
+        if getattr(self.conv3.bn, 'weight', None) is not None:
+            nn.init.zeros_(self.conv3.bn.weight)
+
+    def conv_forward(self, x):
+        conv = self.conv2
+        k = conv.in_channels
+        C = x.size()[1] // k
+        kernel = conv.weight.repeat(C, 1, 1, 1)
+        bias = conv.bias.repeat(C) if conv.bias is not None else None
+        return F.conv2d(x, kernel, bias, conv.stride, 
+            conv.padding, conv.dilation, C * k)
+
+    def forward(self, x):
+        x0 = self.skip(x)
+        x = self.conv1(x)
+        C = x.size(1) // 2
+        x = x[:, :C, :, :] * x[:, C:, :, :]
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.conv3(x)
+        x += x0
+        if self.act3 is not None:
+            x = self.act3(x)
+        return x
+
+def make_blocks(
+        block_fn,
+        channels,
+        block_repeats,
+        inplanes,
+        reduce_first=1,
+        output_stride=32,
+        down_kernel_size=1,
+        avg_down=False,
+        **kwargs,
+):
+    stages = []
+    feature_info = []
+    net_num_blocks = sum(block_repeats)
+    net_block_idx = 0
+    net_stride = 4
+    dilation = prev_dilation = 1
+    for stage_idx, (planes, num_blocks) in enumerate(zip(channels, block_repeats)):
+        stage_name = f'layer{stage_idx + 1}'  # never liked this name, but weight compat requires it
+        stride = 1 if stage_idx == 0 else 2
+        if net_stride >= output_stride:
+            dilation *= stride
+            stride = 1
+        else:
+            net_stride *= stride
+
+        downsample = None
+        if stride != 1 or inplanes != planes * block_fn.expansion:
+            downsample = True 
+
+        block_kwargs = dict(reduce_first=reduce_first, dilation=dilation, **kwargs)
+        blocks = []
+        for block_idx in range(num_blocks):
+            downsample = downsample if block_idx == 0 else None
+            stride = stride if block_idx == 0 else 1
+            blocks.append(block_fn(
+                inplanes, planes, stride, downsample, first_dilation=prev_dilation,
+                **block_kwargs))
+            prev_dilation = dilation
+            inplanes = planes * block_fn.expansion
+            net_block_idx += 1
+
+        stages.append((stage_name, nn.Sequential(*blocks)))
+        feature_info.append(dict(num_chs=inplanes, reduction=net_stride, module=stage_name))
+
+    return stages, feature_info
+
+
+class QLNet(nn.Module):
+    # based on timm code for ResNet / ResNeXt / SE-ResNeXt / SE-Net
+
+    def __init__(
+            self,
+            block=QLBlock,   # new block
+            layers=[3,4,6,3], # as in resnet50
+            num_classes=1000,
+            in_chans=3,
+            output_stride=32,
+            global_pool='avg',
+            cardinality=1,
+            base_width=64,
+            stem_width=64,
+            stem_type='',
+            replace_stem_pool=False,
+            block_reduce_first=1,
+            down_kernel_size=1,
+            avg_down=False,
+            act_layer=nn.ReLU,
+            norm_layer=nn.BatchNorm2d,
+            zero_init_last=True,
+            block_args=None,
+    ):
+        """
+        Args:
+            block (nn.Module): class for the residual block. Options are BasicBlock, Bottleneck.
+            layers (List[int]) : number of layers in each block
+            num_classes (int): number of classification classes (default 1000)
+            in_chans (int): number of input (color) channels. (default 3)
+            output_stride (int): output stride of the network, 32, 16, or 8. (default 32)
+            global_pool (str): Global pooling type. One of 'avg', 'max', 'avgmax', 'catavgmax' (default 'avg')
+            cardinality (int): number of convolution groups for 3x3 conv in Bottleneck. (default 1)
+            base_width (int): bottleneck channels factor. `planes * base_width / 64 * cardinality` (default 64)
+            stem_width (int): number of channels in stem convolutions (default 64)
+            stem_type (str): The type of stem (default ''):
+                * '', default - a single 7x7 conv with a width of stem_width
+                * 'deep' - three 3x3 convolution layers of widths stem_width, stem_width, stem_width * 2
+                * 'deep_tiered' - three 3x3 conv layers of widths stem_width//4 * 3, stem_width, stem_width * 2
+            replace_stem_pool (bool): replace stem max-pooling layer with a 3x3 stride-2 convolution
+            block_reduce_first (int): Reduction factor for first convolution output width of residual blocks,
+                1 for all archs except senets, where 2 (default 1)
+            down_kernel_size (int): kernel size of residual block downsample path,
+                1x1 for most, 3x3 for senets (default: 1)
+            avg_down (bool): use avg pooling for projection skip connection between stages/downsample (default False)
+            act_layer (str, nn.Module): activation layer
+            norm_layer (str, nn.Module): normalization layer
+            zero_init_last (bool): zero-init the last weight in residual path (usually last BN affine weight)
+            block_args (dict): Extra kwargs to pass through to block module
+        """
+        super(QLNet, self).__init__()
+        block_args = block_args or dict()
+        assert output_stride in (8, 16, 32)
+        self.num_classes = num_classes
+        self.grad_checkpointing = False
+        
+        act_layer = get_act_layer(act_layer)
+        norm_layer = get_norm_layer(norm_layer)
+
+        # Stem
+        deep_stem = 'deep' in stem_type
+        inplanes = stem_width * 2 if deep_stem else 64
+        if deep_stem:
+            stem_chs = (stem_width, stem_width)
+            if 'tiered' in stem_type:
+                stem_chs = (3 * (stem_width // 4), stem_width)
+            self.conv1 = nn.Sequential(*[
+                nn.Conv2d(in_chans, stem_chs[0], 3, stride=2, padding=1, bias=False),
+                norm_layer(stem_chs[0]),
+                act_layer(inplace=True),
+                nn.Conv2d(stem_chs[0], stem_chs[1], 3, stride=1, padding=1, bias=False),
+                norm_layer(stem_chs[1]),
+                act_layer(inplace=True),
+                nn.Conv2d(stem_chs[1], inplanes, 3, stride=1, padding=1, bias=False)])
+        else:
+            self.conv1 = nn.Conv2d(in_chans, inplanes, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = norm_layer(inplanes)
+        self.act1 = act_layer(inplace=True)
+        self.feature_info = [dict(num_chs=inplanes, reduction=2, module='act1')]
+
+        # Stem pooling. The name 'maxpool' remains for weight compatibility.
+        if replace_stem_pool:
+            self.maxpool = nn.Sequential(*filter(None, [
+                nn.Conv2d(inplanes, inplanes, 3, stride=2, padding=1, bias=False),
+                norm_layer(inplanes),
+                act_layer(inplace=True)
+            ]))
+        else:
+            self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        # Feature Blocks
+        channels = [64, 128, 256, 512]
+        stage_modules, stage_feature_info = make_blocks(
+            block,
+            channels,
+            layers,
+            inplanes,
+            cardinality=cardinality,
+            base_width=base_width,
+            output_stride=output_stride,
+            reduce_first=block_reduce_first,
+            avg_down=avg_down,
+            down_kernel_size=down_kernel_size,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            **block_args,
+        )
+        for stage in stage_modules:
+            self.add_module(*stage)  # layer1, layer2, etc
+        self.feature_info.extend(stage_feature_info)
+
+        self.act = hardball(radius2=512)
+        # self.act = nn.Hardtanh(max_val=5, min_val=-5, inplace=True)
+        # self.act = nn.ReLU(inplace=True)
+
+        # Head (Pooling and Classifier)
+        self.num_features = 512 * block.expansion
+        self.global_pool, self.fc = create_classifier(self.num_features, self.num_classes, pool_type=global_pool)
+
+        self.init_weights(zero_init_last=zero_init_last)
+
+    @staticmethod
+    def from_pretrained(model_name: str, load_weights=True, **kwargs) -> 'ResNet':
+        entry_fn = model_entrypoint(model_name, 'resnet')
+        return entry_fn(pretrained=not load_weights, **kwargs)
+
+    @torch.jit.ignore
+    def init_weights(self, zero_init_last=True):
+        for n, m in self.named_modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='linear') # 'linear' for non-relu activations
+                # nn.init.xavier_normal_(m.weight)
+        if zero_init_last:
+            for m in self.modules():
+                if hasattr(m, 'zero_init_last'):
+                    m.zero_init_last()
+
+    @torch.jit.ignore
+    def group_matcher(self, coarse=False):
+        matcher = dict(stem=r'^conv1|bn1|maxpool', blocks=r'^layer(\d+)' if coarse else r'^layer(\d+)\.(\d+)')
+        return matcher
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def get_classifier(self, name_only=False):
+        return 'fc' if name_only else self.fc
+
+    def reset_classifier(self, num_classes, global_pool='avg'):
+        self.num_classes = num_classes
+        self.global_pool, self.fc = create_classifier(self.num_features, 99, # self.num_classes, 
+            pool_type=global_pool)
+
+    def forward_features(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+        x = self.maxpool(x)
+
+        if self.grad_checkpointing and not torch.jit.is_scripting():
+            x = checkpoint_seq([self.layer1, self.layer2, self.layer3, self.layer4], x, flatten=True)
+        else:
+            x = self.layer1(x)
+            x = self.layer2(x)
+            x = self.layer3(x)
+            x = self.layer4(x)
+        return x
+
+    def forward_head(self, x, pre_logits: bool = False):
+        x = self.global_pool(x)
+        return x if pre_logits else self.fc(x)
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.act(x)
+        x = self.forward_head(x)
+        return x
+