data_exp2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset,DataLoader\n",
    "from glob import glob\n",
    "from sklearn.model_selection import train_test_split\n",
    "import ipdb\n",
    "import rasterio\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "from tqdm.auto import tqdm\n",
    "from model import WrapperModel\n",
    "from dataset import SARDataset\n",
    "import segmentation_models_pytorch as smp\n",
    "import torch.nn.functional as F\n",
    "import pathlib\n",
    "import os\n",
    "from torch import nn\n",
    "import albumentations as A\n",
    "import ttach as tta\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "WrapperModel(\n",
       "  (model): UnetPlusPlus(\n",
       "    (encoder): TimmUniversalEncoder(\n",
       "      (model): FeatureListNet(\n",
       "        (stem): Stem(\n",
       "          (conv1): Conv2dSame(6, 128, kernel_size=(3, 3), stride=(2, 2))\n",
       "          (norm1): BatchNormAct2d(\n",
       "            128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "            (drop): Identity()\n",
       "            (act): GELUTanh()\n",
       "          )\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "        (stages_0): MaxxVitStage(\n",
       "          (blocks): Sequential(\n",
       "            (0): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Downsample2d(\n",
       "                  (pool): AvgPool2dSame(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0))\n",
       "                  (expand): Identity()\n",
       "                )\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2dSame(512, 512, kernel_size=(3, 3), stride=(2, 2), groups=512, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=128, out_features=384, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=128, out_features=384, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (1): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=128, out_features=384, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=128, out_features=384, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (stages_1): MaxxVitStage(\n",
       "          (blocks): Sequential(\n",
       "            (0): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Downsample2d(\n",
       "                  (pool): AvgPool2dSame(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0))\n",
       "                  (expand): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "                )\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(128, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2dSame(1024, 1024, kernel_size=(3, 3), stride=(2, 2), groups=1024, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (1): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (2): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (3): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (4): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (5): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (stages_2): MaxxVitStage(\n",
       "          (blocks): Sequential(\n",
       "            (0): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Downsample2d(\n",
       "                  (pool): AvgPool2dSame(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0))\n",
       "                  (expand): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                )\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(256, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2dSame(2048, 2048, kernel_size=(3, 3), stride=(2, 2), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (1): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (2): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (3): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (4): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (5): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (6): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (7): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (8): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (9): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (10): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (11): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (12): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (13): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2048, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  2048, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=512, out_features=1536, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=512, out_features=512, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (stages_3): MaxxVitStage(\n",
       "          (blocks): Sequential(\n",
       "            (0): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Downsample2d(\n",
       "                  (pool): AvgPool2dSame(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0))\n",
       "                  (expand): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "                )\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(512, 4096, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  4096, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2dSame(4096, 4096, kernel_size=(3, 3), stride=(2, 2), groups=4096, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  4096, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(4096, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(256, 4096, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(4096, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "            (1): MaxxVitBlock(\n",
       "              (conv): MbConvBlock(\n",
       "                (shortcut): Identity()\n",
       "                (pre_norm): BatchNormAct2d(\n",
       "                  1024, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): Identity()\n",
       "                )\n",
       "                (down): Identity()\n",
       "                (conv1_1x1): Conv2d(1024, 4096, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "                (norm1): BatchNormAct2d(\n",
       "                  4096, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (conv2_kxk): Conv2d(4096, 4096, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=4096, bias=False)\n",
       "                (norm2): BatchNormAct2d(\n",
       "                  4096, eps=0.001, momentum=0.1, affine=True, track_running_stats=True\n",
       "                  (drop): Identity()\n",
       "                  (act): GELUTanh()\n",
       "                )\n",
       "                (se): SEModule(\n",
       "                  (fc1): Conv2d(4096, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (bn): Identity()\n",
       "                  (act): SiLU(inplace=True)\n",
       "                  (fc2): Conv2d(256, 4096, kernel_size=(1, 1), stride=(1, 1))\n",
       "                  (gate): Sigmoid()\n",
       "                )\n",
       "                (conv3_1x1): Conv2d(4096, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "                (drop_path): Identity()\n",
       "              )\n",
       "              (attn_block): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "              (attn_grid): PartitionAttentionCl(\n",
       "                (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "                (attn): AttentionCl(\n",
       "                  (qkv): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "                  (rel_pos): RelPosBiasTf()\n",
       "                  (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                  (proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "                  (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls1): Identity()\n",
       "                (drop_path1): Identity()\n",
       "                (norm2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "                (mlp): Mlp(\n",
       "                  (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (act): GELUTanh()\n",
       "                  (drop1): Dropout(p=0.0, inplace=False)\n",
       "                  (norm): Identity()\n",
       "                  (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "                  (drop2): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (ls2): Identity()\n",
       "                (drop_path2): Identity()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (decoder): UnetPlusPlusDecoder(\n",
       "      (center): Identity()\n",
       "      (blocks): ModuleDict(\n",
       "        (x_0_0): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "        (x_0_1): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(768, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "        (x_1_1): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(768, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "        (x_0_2): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(512, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "        (x_1_2): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(512, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "        (x_2_2): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(384, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "        (x_0_3): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(576, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "        (x_1_3): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(512, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "        (x_2_3): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(384, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "        (x_3_3): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "        (x_0_4): DecoderBlock(\n",
       "          (conv1): Conv2dReLU(\n",
       "            (0): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention1): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "          (conv2): Conv2dReLU(\n",
       "            (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (attention2): Attention(\n",
       "            (attention): Identity()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (segmentation_head): SegmentationHead(\n",
       "      (0): Conv2d(16, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): Identity()\n",
       "      (2): Activation(\n",
       "        (activation): Identity()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (bce_logit_loss): BCEWithLogitsLoss()\n",
       "  (dice_loss): DiceLoss()\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unet_pp=smp.create_model(arch='unetplusplus',classes=1,in_channels=6,encoder_name='tu-maxvit_large_tf_512', encoder_weights=None)\n",
    "model1=WrapperModel.load_from_checkpoint('/home/syo/work/2024_IEEE_GRSS/src/weights/maxvit-800e/epoch=791-step=72864.ckpt',model=unet_pp,mode='test',map_location=torch.device(\"cuda\"))\n",
    "model1.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _sar_normalization(sar_data):\n",
    "    \"\"\"\n",
    "        normalize the data to (0, 1).\n",
    "\n",
    "        Band 1: SAR image, VV\n",
    "        Band 2: SAR image, VH\n",
    "        Band 3: Merit DEM\n",
    "        Band 4: Copernicus DEM\n",
    "        Band 5: ESA World Cover Map\n",
    "        Band 6: Water occurrence probability\n",
    "\n",
    "        # 1: 0-32765        x/32765\n",
    "        # 2:                x/32765\n",
    "        # 3: -9999:9999     (x+9999)/(9999*2)\n",
    "        # 4: 0-100          x/100\n",
    "        # 5: 0-255          x/255\n",
    "    \"\"\"\n",
    "    sar_data=sar_data.astype(np.float32)\n",
    "    sar_data[0]=sar_data[0]/32765\n",
    "    sar_data[1]=sar_data[1]/32765\n",
    "    sar_data[2]=(sar_data[2]+9999)/(9999*2)\n",
    "    sar_data[3]=(sar_data[3]+9999)/(9999*2)\n",
    "    sar_data[4]=sar_data[4]/100\n",
    "    sar_data[5]=sar_data[5]/255\n",
    "\n",
    "    return sar_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/349 [00:00<?, ?it/s]/home/syo/opt/mambaforge/envs/kaggle/lib/python3.11/site-packages/rasterio/__init__.py:304: NotGeoreferencedWarning: Dataset has no geotransform, gcps, or rpcs. The identity matrix will be returned.\n",
      "  dataset = DatasetReader(path, driver=driver, sharing=sharing, **kwargs)\n",
      "100%|██████████| 349/349 [00:28<00:00, 12.31it/s]\n"
     ]
    }
   ],
   "source": [
    "image_root=pathlib.Path('/home/syo/work/2024_IEEE_GRSS/dataset/Track1/val/images/')\n",
    "image_list = sorted(list(image_root.glob('*')))\n",
    "thred=1e-3\n",
    "\n",
    "num_rows = len(image_list) // 3 + 1\n",
    "num_cols = 3\n",
    "\n",
    "plt.figure(figsize=(num_rows,num_cols))\n",
    "\n",
    "for n,i in tqdm(enumerate(image_list, 1), total=len(image_list)):\n",
    "    sar_data = rasterio.open(i).read()\n",
    "    sar_data=_sar_normalization(sar_data=sar_data)\n",
    "    with torch.no_grad():\n",
    "            sar_data=torch.tensor(sar_data,dtype=torch.float32).unsqueeze(0).cuda()\n",
    "            outputs1=model1.forward(sar_data)\n",
    "            outputs1 = F.sigmoid(outputs1)\n",
    "            outputs=outputs1\n",
    "            preds=outputs.detach().cpu().numpy().squeeze(0)[0]\n",
    "            binary_map = (preds > thred).astype(np.uint8)\n",
    "            #plt.imshow(preds)\n",
    "            plt.subplot(num_rows,num_cols,n)\n",
    "            plt.imshow(binary_map)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "kaggle",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}