Segmentation working!

1 files changed, 252 insertions, 173 deletions

diff --git a/src/notebooks/00-testing-stuff-out.ipynb b/src/notebooks/00-testing-stuff-out.ipynb
index 96a0c5a..6f01dfb 100644
--- a/src/notebooks/00-testing-stuff-out.ipynb
+++ b/src/notebooks/00-testing-stuff-out.ipynb
@@ -2,9 +2,18 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 21,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The autoreload extension is already loaded. To reload it, use:\n",
+      "  %reload_ext autoreload\n"
+     ]
+    }
+   ],
    "source": [
     "%load_ext autoreload\n",
     "%autoreload 2\n",
@@ -23,16 +32,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from text_recognizer.networks.residual_network import IdentityBlock, ResidualBlock, BasicBlock, BottleNeckBlock, ResidualLayer, ResidualNetwork"
+    "from text_recognizer.networks.residual_network import IdentityBlock, ResidualBlock, BasicBlock, BottleNeckBlock, ResidualLayer, ResidualNetwork, ResidualNetworkEncoder"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -41,7 +50,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -50,7 +59,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -59,7 +68,7 @@
        "True"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -70,7 +79,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -79,7 +88,7 @@
        "False"
       ]
      },
-     "execution_count": 4,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -90,26 +99,26 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 23,
    "metadata": {},
    "outputs": [],
    "source": [
     "wr = WideResidualNetwork(\n",
     "            in_channels= 1,\n",
     "            num_classes= 80,\n",
-    "            in_planes=32,\n",
+    "            in_planes=64,\n",
     "            depth=10,\n",
     "            num_layers=4,\n",
     "            width_factor=1,\n",
     "            dropout_rate= 0.2,\n",
     "            activation= \"SELU\",\n",
-    "            use_decoder= False,\n",
+    "            use_decoder= True,\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -118,18 +127,102 @@
   },
   {
    "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "backbone = ResidualNetworkEncoder(1, [64, 128, 256], [2, 2, 3])"
+   ]
+  },
+  {
+   "cell_type": "code",
    "execution_count": 20,
    "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "==========================================================================================\n",
+      "Layer (type:depth-idx)                   Output Shape              Param #\n",
+      "==========================================================================================\n",
+      "├─Sequential: 1-1                        [-1, 64, 28, 952]         --\n",
+      "|    └─Conv2d: 2-1                       [-1, 64, 28, 952]         576\n",
+      "|    └─BatchNorm2d: 2-2                  [-1, 64, 28, 952]         128\n",
+      "|    └─ReLU: 2-3                         [-1, 64, 28, 952]         --\n",
+      "├─Sequential: 1-2                        [-1, 256, 7, 238]         --\n",
+      "|    └─ResidualLayer: 2-4                [-1, 64, 28, 952]         --\n",
+      "|    |    └─Sequential: 3-1              [-1, 64, 28, 952]         147,968\n",
+      "|    └─ResidualLayer: 2-5                [-1, 128, 14, 476]        --\n",
+      "|    |    └─Sequential: 3-2              [-1, 128, 14, 476]        525,568\n",
+      "|    └─ResidualLayer: 2-6                [-1, 256, 7, 238]         --\n",
+      "|    |    └─Sequential: 3-3              [-1, 256, 7, 238]         3,280,384\n",
+      "==========================================================================================\n",
+      "Total params: 3,954,624\n",
+      "Trainable params: 3,954,624\n",
+      "Non-trainable params: 0\n",
+      "Total mult-adds (M): 31.16\n",
+      "==========================================================================================\n",
+      "Input size (MB): 0.10\n",
+      "Forward/backward pass size (MB): 26.03\n",
+      "Params size (MB): 15.09\n",
+      "Estimated Total Size (MB): 41.22\n",
+      "==========================================================================================\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "==========================================================================================\n",
+       "Layer (type:depth-idx)                   Output Shape              Param #\n",
+       "==========================================================================================\n",
+       "├─Sequential: 1-1                        [-1, 64, 28, 952]         --\n",
+       "|    └─Conv2d: 2-1                       [-1, 64, 28, 952]         576\n",
+       "|    └─BatchNorm2d: 2-2                  [-1, 64, 28, 952]         128\n",
+       "|    └─ReLU: 2-3                         [-1, 64, 28, 952]         --\n",
+       "├─Sequential: 1-2                        [-1, 256, 7, 238]         --\n",
+       "|    └─ResidualLayer: 2-4                [-1, 64, 28, 952]         --\n",
+       "|    |    └─Sequential: 3-1              [-1, 64, 28, 952]         147,968\n",
+       "|    └─ResidualLayer: 2-5                [-1, 128, 14, 476]        --\n",
+       "|    |    └─Sequential: 3-2              [-1, 128, 14, 476]        525,568\n",
+       "|    └─ResidualLayer: 2-6                [-1, 256, 7, 238]         --\n",
+       "|    |    └─Sequential: 3-3              [-1, 256, 7, 238]         3,280,384\n",
+       "==========================================================================================\n",
+       "Total params: 3,954,624\n",
+       "Trainable params: 3,954,624\n",
+       "Non-trainable params: 0\n",
+       "Total mult-adds (M): 31.16\n",
+       "==========================================================================================\n",
+       "Input size (MB): 0.10\n",
+       "Forward/backward pass size (MB): 26.03\n",
+       "Params size (MB): 15.09\n",
+       "Estimated Total Size (MB): 41.22\n",
+       "=========================================================================================="
+      ]
+     },
+     "execution_count": 20,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "summary(backbone, (1, 28, 952), device=\"cpu\", depth=3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
    "outputs": [],
    "source": [
     "        backbone = nn.Sequential(\n",
-    "            *list(wr.children())[:][:-1]\n",
+    "            *list(wr.children())[:][:]\n",
     "        )\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 25,
    "metadata": {},
    "outputs": [
     {
@@ -137,10 +230,84 @@
       "text/plain": [
        "Sequential(\n",
        "  (0): SELU(inplace=True)\n",
+       "  (1): Sequential(\n",
+       "    (0): Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+       "    (1): Sequential(\n",
+       "      (0): WideBlock(\n",
+       "        (activation): SELU(inplace=True)\n",
+       "        (blocks): Sequential(\n",
+       "          (0): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (1): SELU(inplace=True)\n",
+       "          (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+       "          (3): Dropout(p=0.2, inplace=False)\n",
+       "          (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (5): SELU(inplace=True)\n",
+       "          (6): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+       "        )\n",
+       "      )\n",
+       "    )\n",
+       "    (2): Sequential(\n",
+       "      (0): WideBlock(\n",
+       "        (activation): SELU(inplace=True)\n",
+       "        (blocks): Sequential(\n",
+       "          (0): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (1): SELU(inplace=True)\n",
+       "          (2): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+       "          (3): Dropout(p=0.2, inplace=False)\n",
+       "          (4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (5): SELU(inplace=True)\n",
+       "          (6): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
+       "        )\n",
+       "        (shortcut): Sequential(\n",
+       "          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
+       "        )\n",
+       "      )\n",
+       "    )\n",
+       "    (3): Sequential(\n",
+       "      (0): WideBlock(\n",
+       "        (activation): SELU(inplace=True)\n",
+       "        (blocks): Sequential(\n",
+       "          (0): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (1): SELU(inplace=True)\n",
+       "          (2): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+       "          (3): Dropout(p=0.2, inplace=False)\n",
+       "          (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (5): SELU(inplace=True)\n",
+       "          (6): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
+       "        )\n",
+       "        (shortcut): Sequential(\n",
+       "          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
+       "        )\n",
+       "      )\n",
+       "    )\n",
+       "    (4): Sequential(\n",
+       "      (0): WideBlock(\n",
+       "        (activation): SELU(inplace=True)\n",
+       "        (blocks): Sequential(\n",
+       "          (0): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (1): SELU(inplace=True)\n",
+       "          (2): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+       "          (3): Dropout(p=0.2, inplace=False)\n",
+       "          (4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (5): SELU(inplace=True)\n",
+       "          (6): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
+       "        )\n",
+       "        (shortcut): Sequential(\n",
+       "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
+       "        )\n",
+       "      )\n",
+       "    )\n",
+       "  )\n",
+       "  (2): Sequential(\n",
+       "    (0): BatchNorm2d(512, eps=1e-05, momentum=0.8, affine=True, track_running_stats=True)\n",
+       "    (1): SELU(inplace=True)\n",
+       "    (2): Reduce('b c h w -> b c', 'mean')\n",
+       "    (3): Linear(in_features=512, out_features=80, bias=True)\n",
+       "  )\n",
        ")"
       ]
      },
-     "execution_count": 21,
+     "execution_count": 25,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -151,188 +318,100 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 86,
+   "execution_count": 26,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "===============================================================================================\n",
-      "Layer (type:depth-idx)                        Output Shape              Param #\n",
-      "===============================================================================================\n",
-      "├─Sequential: 1-1                             [-1, 256, 4, 119]         --\n",
-      "|    └─Conv2d: 2-1                            [-1, 32, 28, 952]         288\n",
-      "|    └─Sequential: 2-2                        [-1, 32, 28, 952]         --\n",
-      "|    |    └─WideBlock: 3-1                    [-1, 32, 28, 952]         --\n",
-      "|    |    |    └─Sequential: 4-1              [-1, 32, 28, 952]         --\n",
-      "|    |    |    |    └─BatchNorm2d: 5-1        [-1, 32, 28, 952]         64\n",
-      "|    |    |    └─SELU: 4-2                    [-1, 32, 28, 952]         --\n",
-      "|    |    |    └─Sequential: 4                []                        --\n",
-      "|    |    |    |    └─SELU: 5-2               [-1, 32, 28, 952]         --\n",
-      "|    |    |    |    └─Conv2d: 5-3             [-1, 32, 28, 952]         9,216\n",
-      "|    |    |    |    └─Dropout: 5-4            [-1, 32, 28, 952]         --\n",
-      "|    |    |    |    └─BatchNorm2d: 5-5        [-1, 32, 28, 952]         64\n",
-      "|    |    |    └─SELU: 4-3                    [-1, 32, 28, 952]         --\n",
-      "|    |    |    └─Sequential: 4                []                        --\n",
-      "|    |    |    |    └─SELU: 5-6               [-1, 32, 28, 952]         --\n",
-      "|    |    |    |    └─Conv2d: 5-7             [-1, 32, 28, 952]         9,216\n",
-      "|    └─Sequential: 2-3                        [-1, 64, 14, 476]         --\n",
-      "|    |    └─WideBlock: 3-2                    [-1, 64, 14, 476]         --\n",
-      "|    |    |    └─Sequential: 4-4              [-1, 64, 14, 476]         --\n",
-      "|    |    |    |    └─Conv2d: 5-8             [-1, 64, 14, 476]         2,048\n",
-      "|    |    |    └─Sequential: 4-5              [-1, 64, 14, 476]         --\n",
-      "|    |    |    |    └─BatchNorm2d: 5-9        [-1, 32, 28, 952]         64\n",
-      "|    |    |    └─SELU: 4-6                    [-1, 32, 28, 952]         --\n",
-      "|    |    |    └─Sequential: 4                []                        --\n",
-      "|    |    |    |    └─SELU: 5-10              [-1, 32, 28, 952]         --\n",
-      "|    |    |    |    └─Conv2d: 5-11            [-1, 64, 28, 952]         18,432\n",
-      "|    |    |    |    └─Dropout: 5-12           [-1, 64, 28, 952]         --\n",
-      "|    |    |    |    └─BatchNorm2d: 5-13       [-1, 64, 28, 952]         128\n",
-      "|    |    |    └─SELU: 4-7                    [-1, 64, 28, 952]         --\n",
-      "|    |    |    └─Sequential: 4                []                        --\n",
-      "|    |    |    |    └─SELU: 5-14              [-1, 64, 28, 952]         --\n",
-      "|    |    |    |    └─Conv2d: 5-15            [-1, 64, 14, 476]         36,864\n",
-      "|    └─Sequential: 2-4                        [-1, 128, 7, 238]         --\n",
-      "|    |    └─WideBlock: 3-3                    [-1, 128, 7, 238]         --\n",
-      "|    |    |    └─Sequential: 4-8              [-1, 128, 7, 238]         --\n",
-      "|    |    |    |    └─Conv2d: 5-16            [-1, 128, 7, 238]         8,192\n",
-      "|    |    |    └─Sequential: 4-9              [-1, 128, 7, 238]         --\n",
-      "|    |    |    |    └─BatchNorm2d: 5-17       [-1, 64, 14, 476]         128\n",
-      "|    |    |    └─SELU: 4-10                   [-1, 64, 14, 476]         --\n",
-      "|    |    |    └─Sequential: 4                []                        --\n",
-      "|    |    |    |    └─SELU: 5-18              [-1, 64, 14, 476]         --\n",
-      "|    |    |    |    └─Conv2d: 5-19            [-1, 128, 14, 476]        73,728\n",
-      "|    |    |    |    └─Dropout: 5-20           [-1, 128, 14, 476]        --\n",
-      "|    |    |    |    └─BatchNorm2d: 5-21       [-1, 128, 14, 476]        256\n",
-      "|    |    |    └─SELU: 4-11                   [-1, 128, 14, 476]        --\n",
-      "|    |    |    └─Sequential: 4                []                        --\n",
-      "|    |    |    |    └─SELU: 5-22              [-1, 128, 14, 476]        --\n",
-      "|    |    |    |    └─Conv2d: 5-23            [-1, 128, 7, 238]         147,456\n",
-      "|    └─Sequential: 2-5                        [-1, 256, 4, 119]         --\n",
-      "|    |    └─WideBlock: 3-4                    [-1, 256, 4, 119]         --\n",
-      "|    |    |    └─Sequential: 4-12             [-1, 256, 4, 119]         --\n",
-      "|    |    |    |    └─Conv2d: 5-24            [-1, 256, 4, 119]         32,768\n",
-      "|    |    |    └─Sequential: 4-13             [-1, 256, 4, 119]         --\n",
-      "|    |    |    |    └─BatchNorm2d: 5-25       [-1, 128, 7, 238]         256\n",
-      "|    |    |    └─SELU: 4-14                   [-1, 128, 7, 238]         --\n",
-      "|    |    |    └─Sequential: 4                []                        --\n",
-      "|    |    |    |    └─SELU: 5-26              [-1, 128, 7, 238]         --\n",
-      "|    |    |    |    └─Conv2d: 5-27            [-1, 256, 7, 238]         294,912\n",
-      "|    |    |    |    └─Dropout: 5-28           [-1, 256, 7, 238]         --\n",
-      "|    |    |    |    └─BatchNorm2d: 5-29       [-1, 256, 7, 238]         512\n",
-      "|    |    |    └─SELU: 4-15                   [-1, 256, 7, 238]         --\n",
-      "|    |    |    └─Sequential: 4                []                        --\n",
-      "|    |    |    |    └─SELU: 5-30              [-1, 256, 7, 238]         --\n",
-      "|    |    |    |    └─Conv2d: 5-31            [-1, 256, 4, 119]         589,824\n",
-      "===============================================================================================\n",
-      "Total params: 1,224,416\n",
-      "Trainable params: 1,224,416\n",
+      "==========================================================================================\n",
+      "Layer (type:depth-idx)                   Output Shape              Param #\n",
+      "==========================================================================================\n",
+      "├─SELU: 1-1                              [-1, 1, 28, 952]          --\n",
+      "├─Sequential: 1                          []                        --\n",
+      "|    └─SELU: 2-1                         [-1, 1, 28, 952]          --\n",
+      "├─Sequential: 1-2                        [-1, 512, 4, 119]         --\n",
+      "|    └─Conv2d: 2-2                       [-1, 64, 28, 952]         576\n",
+      "|    └─Sequential: 2-3                   [-1, 64, 28, 952]         --\n",
+      "|    |    └─WideBlock: 3-1               [-1, 64, 28, 952]         73,984\n",
+      "|    └─Sequential: 2-4                   [-1, 128, 14, 476]        --\n",
+      "|    |    └─WideBlock: 3-2               [-1, 128, 14, 476]        229,760\n",
+      "|    └─Sequential: 2-5                   [-1, 256, 7, 238]         --\n",
+      "|    |    └─WideBlock: 3-3               [-1, 256, 7, 238]         918,272\n",
+      "|    └─Sequential: 2-6                   [-1, 512, 4, 119]         --\n",
+      "|    |    └─WideBlock: 3-4               [-1, 512, 4, 119]         3,671,552\n",
+      "├─Sequential: 1-3                        [-1, 80]                  --\n",
+      "|    └─BatchNorm2d: 2-7                  [-1, 512, 4, 119]         1,024\n",
+      "├─SELU: 1-4                              [-1, 512, 4, 119]         --\n",
+      "├─Sequential: 1                          []                        --\n",
+      "|    └─SELU: 2-8                         [-1, 512, 4, 119]         --\n",
+      "|    └─Reduce: 2-9                       [-1, 512]                 --\n",
+      "|    └─Linear: 2-10                      [-1, 80]                  41,040\n",
+      "==========================================================================================\n",
+      "Total params: 4,936,208\n",
+      "Trainable params: 4,936,208\n",
       "Non-trainable params: 0\n",
-      "Total mult-adds (G): 2.79\n",
-      "===============================================================================================\n",
+      "Total mult-adds (M): 35.01\n",
+      "==========================================================================================\n",
       "Input size (MB): 0.10\n",
-      "Forward/backward pass size (MB): 101.10\n",
-      "Params size (MB): 4.67\n",
-      "Estimated Total Size (MB): 105.88\n",
-      "===============================================================================================\n"
+      "Forward/backward pass size (MB): 14.88\n",
+      "Params size (MB): 18.83\n",
+      "Estimated Total Size (MB): 33.81\n",
+      "==========================================================================================\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "===============================================================================================\n",
-       "Layer (type:depth-idx)                        Output Shape              Param #\n",
-       "===============================================================================================\n",
-       "├─Sequential: 1-1                             [-1, 256, 4, 119]         --\n",
-       "|    └─Conv2d: 2-1                            [-1, 32, 28, 952]         288\n",
-       "|    └─Sequential: 2-2                        [-1, 32, 28, 952]         --\n",
-       "|    |    └─WideBlock: 3-1                    [-1, 32, 28, 952]         --\n",
-       "|    |    |    └─Sequential: 4-1              [-1, 32, 28, 952]         --\n",
-       "|    |    |    |    └─BatchNorm2d: 5-1        [-1, 32, 28, 952]         64\n",
-       "|    |    |    └─SELU: 4-2                    [-1, 32, 28, 952]         --\n",
-       "|    |    |    └─Sequential: 4                []                        --\n",
-       "|    |    |    |    └─SELU: 5-2               [-1, 32, 28, 952]         --\n",
-       "|    |    |    |    └─Conv2d: 5-3             [-1, 32, 28, 952]         9,216\n",
-       "|    |    |    |    └─Dropout: 5-4            [-1, 32, 28, 952]         --\n",
-       "|    |    |    |    └─BatchNorm2d: 5-5        [-1, 32, 28, 952]         64\n",
-       "|    |    |    └─SELU: 4-3                    [-1, 32, 28, 952]         --\n",
-       "|    |    |    └─Sequential: 4                []                        --\n",
-       "|    |    |    |    └─SELU: 5-6               [-1, 32, 28, 952]         --\n",
-       "|    |    |    |    └─Conv2d: 5-7             [-1, 32, 28, 952]         9,216\n",
-       "|    └─Sequential: 2-3                        [-1, 64, 14, 476]         --\n",
-       "|    |    └─WideBlock: 3-2                    [-1, 64, 14, 476]         --\n",
-       "|    |    |    └─Sequential: 4-4              [-1, 64, 14, 476]         --\n",
-       "|    |    |    |    └─Conv2d: 5-8             [-1, 64, 14, 476]         2,048\n",
-       "|    |    |    └─Sequential: 4-5              [-1, 64, 14, 476]         --\n",
-       "|    |    |    |    └─BatchNorm2d: 5-9        [-1, 32, 28, 952]         64\n",
-       "|    |    |    └─SELU: 4-6                    [-1, 32, 28, 952]         --\n",
-       "|    |    |    └─Sequential: 4                []                        --\n",
-       "|    |    |    |    └─SELU: 5-10              [-1, 32, 28, 952]         --\n",
-       "|    |    |    |    └─Conv2d: 5-11            [-1, 64, 28, 952]         18,432\n",
-       "|    |    |    |    └─Dropout: 5-12           [-1, 64, 28, 952]         --\n",
-       "|    |    |    |    └─BatchNorm2d: 5-13       [-1, 64, 28, 952]         128\n",
-       "|    |    |    └─SELU: 4-7                    [-1, 64, 28, 952]         --\n",
-       "|    |    |    └─Sequential: 4                []                        --\n",
-       "|    |    |    |    └─SELU: 5-14              [-1, 64, 28, 952]         --\n",
-       "|    |    |    |    └─Conv2d: 5-15            [-1, 64, 14, 476]         36,864\n",
-       "|    └─Sequential: 2-4                        [-1, 128, 7, 238]         --\n",
-       "|    |    └─WideBlock: 3-3                    [-1, 128, 7, 238]         --\n",
-       "|    |    |    └─Sequential: 4-8              [-1, 128, 7, 238]         --\n",
-       "|    |    |    |    └─Conv2d: 5-16            [-1, 128, 7, 238]         8,192\n",
-       "|    |    |    └─Sequential: 4-9              [-1, 128, 7, 238]         --\n",
-       "|    |    |    |    └─BatchNorm2d: 5-17       [-1, 64, 14, 476]         128\n",
-       "|    |    |    └─SELU: 4-10                   [-1, 64, 14, 476]         --\n",
-       "|    |    |    └─Sequential: 4                []                        --\n",
-       "|    |    |    |    └─SELU: 5-18              [-1, 64, 14, 476]         --\n",
-       "|    |    |    |    └─Conv2d: 5-19            [-1, 128, 14, 476]        73,728\n",
-       "|    |    |    |    └─Dropout: 5-20           [-1, 128, 14, 476]        --\n",
-       "|    |    |    |    └─BatchNorm2d: 5-21       [-1, 128, 14, 476]        256\n",
-       "|    |    |    └─SELU: 4-11                   [-1, 128, 14, 476]        --\n",
-       "|    |    |    └─Sequential: 4                []                        --\n",
-       "|    |    |    |    └─SELU: 5-22              [-1, 128, 14, 476]        --\n",
-       "|    |    |    |    └─Conv2d: 5-23            [-1, 128, 7, 238]         147,456\n",
-       "|    └─Sequential: 2-5                        [-1, 256, 4, 119]         --\n",
-       "|    |    └─WideBlock: 3-4                    [-1, 256, 4, 119]         --\n",
-       "|    |    |    └─Sequential: 4-12             [-1, 256, 4, 119]         --\n",
-       "|    |    |    |    └─Conv2d: 5-24            [-1, 256, 4, 119]         32,768\n",
-       "|    |    |    └─Sequential: 4-13             [-1, 256, 4, 119]         --\n",
-       "|    |    |    |    └─BatchNorm2d: 5-25       [-1, 128, 7, 238]         256\n",
-       "|    |    |    └─SELU: 4-14                   [-1, 128, 7, 238]         --\n",
-       "|    |    |    └─Sequential: 4                []                        --\n",
-       "|    |    |    |    └─SELU: 5-26              [-1, 128, 7, 238]         --\n",
-       "|    |    |    |    └─Conv2d: 5-27            [-1, 256, 7, 238]         294,912\n",
-       "|    |    |    |    └─Dropout: 5-28           [-1, 256, 7, 238]         --\n",
-       "|    |    |    |    └─BatchNorm2d: 5-29       [-1, 256, 7, 238]         512\n",
-       "|    |    |    └─SELU: 4-15                   [-1, 256, 7, 238]         --\n",
-       "|    |    |    └─Sequential: 4                []                        --\n",
-       "|    |    |    |    └─SELU: 5-30              [-1, 256, 7, 238]         --\n",
-       "|    |    |    |    └─Conv2d: 5-31            [-1, 256, 4, 119]         589,824\n",
-       "===============================================================================================\n",
-       "Total params: 1,224,416\n",
-       "Trainable params: 1,224,416\n",
+       "==========================================================================================\n",
+       "Layer (type:depth-idx)                   Output Shape              Param #\n",
+       "==========================================================================================\n",
+       "├─SELU: 1-1                              [-1, 1, 28, 952]          --\n",
+       "├─Sequential: 1                          []                        --\n",
+       "|    └─SELU: 2-1                         [-1, 1, 28, 952]          --\n",
+       "├─Sequential: 1-2                        [-1, 512, 4, 119]         --\n",
+       "|    └─Conv2d: 2-2                       [-1, 64, 28, 952]         576\n",
+       "|    └─Sequential: 2-3                   [-1, 64, 28, 952]         --\n",
+       "|    |    └─WideBlock: 3-1               [-1, 64, 28, 952]         73,984\n",
+       "|    └─Sequential: 2-4                   [-1, 128, 14, 476]        --\n",
+       "|    |    └─WideBlock: 3-2               [-1, 128, 14, 476]        229,760\n",
+       "|    └─Sequential: 2-5                   [-1, 256, 7, 238]         --\n",
+       "|    |    └─WideBlock: 3-3               [-1, 256, 7, 238]         918,272\n",
+       "|    └─Sequential: 2-6                   [-1, 512, 4, 119]         --\n",
+       "|    |    └─WideBlock: 3-4               [-1, 512, 4, 119]         3,671,552\n",
+       "├─Sequential: 1-3                        [-1, 80]                  --\n",
+       "|    └─BatchNorm2d: 2-7                  [-1, 512, 4, 119]         1,024\n",
+       "├─SELU: 1-4                              [-1, 512, 4, 119]         --\n",
+       "├─Sequential: 1                          []                        --\n",
+       "|    └─SELU: 2-8                         [-1, 512, 4, 119]         --\n",
+       "|    └─Reduce: 2-9                       [-1, 512]                 --\n",
+       "|    └─Linear: 2-10                      [-1, 80]                  41,040\n",
+       "==========================================================================================\n",
+       "Total params: 4,936,208\n",
+       "Trainable params: 4,936,208\n",
        "Non-trainable params: 0\n",
-       "Total mult-adds (G): 2.79\n",
-       "===============================================================================================\n",
+       "Total mult-adds (M): 35.01\n",
+       "==========================================================================================\n",
        "Input size (MB): 0.10\n",
-       "Forward/backward pass size (MB): 101.10\n",
-       "Params size (MB): 4.67\n",
-       "Estimated Total Size (MB): 105.88\n",
-       "==============================================================================================="
+       "Forward/backward pass size (MB): 14.88\n",
+       "Params size (MB): 18.83\n",
+       "Estimated Total Size (MB): 33.81\n",
+       "=========================================================================================="
       ]
      },
-     "execution_count": 86,
+     "execution_count": 26,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "summary(wr, (1, 28, 952), device=\"cpu\", depth=7)"
+    "summary(backbone, (1, 28, 952), device=\"cpu\", depth=3)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 27,
    "metadata": {},
    "outputs": [],
    "source": [