Model_explore.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/arpitsah/anaconda3/envs/emavfi/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append('emavfi')\n",
    "from emavfi.utils import InputPadder\n",
    "from emavfi.Trainer import Model\n",
    "import emavfi.config as cfg\n",
    "import cv2\n",
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import argparse\n",
    "from glob import glob\n",
    "from natsort import natsorted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/arpitsah/anaconda3/envs/emavfi/lib/python3.8/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/home/arpitsah/anaconda3/envs/emavfi/lib/python3.8/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MultiScaleFlow(\n",
      "  (feature_bone): MotionFormer(\n",
      "    (block1): ConvBlock(\n",
      "      (conv): Sequential(\n",
      "        (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=16)\n",
      "        (2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (3): PReLU(num_parameters=16)\n",
      "      )\n",
      "    )\n",
      "    (patch_embed2): Sequential(\n",
      "      (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=32)\n",
      "    )\n",
      "    (block2): ConvBlock(\n",
      "      (conv): Sequential(\n",
      "        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=32)\n",
      "        (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (3): PReLU(num_parameters=32)\n",
      "      )\n",
      "    )\n",
      "    (patch_embed3): Sequential(\n",
      "      (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (block3): ConvBlock(\n",
      "      (conv): Sequential(\n",
      "        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "        (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (3): PReLU(num_parameters=64)\n",
      "      )\n",
      "    )\n",
      "    (norm4): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "    (patch_embed4): CrossScalePatchEmbed(\n",
      "      (layers): ModuleList(\n",
      "        (0): Conv2d(64, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(1, 1))\n",
      "        (2): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(2, 2), dilation=(2, 2))\n",
      "        (3): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(1, 1))\n",
      "        (4): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(2, 2), dilation=(2, 2))\n",
      "        (5): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(3, 3), dilation=(3, 3))\n",
      "        (6): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(4, 4), dilation=(4, 4))\n",
      "      )\n",
      "      (proj): Conv2d(112, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "      (norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "    (block4): ModuleList(\n",
      "      (0-1): 2 x MotionFormerBlock(\n",
      "        (norm1): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "        (attn): InterFrameAttention(\n",
      "          (q): Linear(in_features=128, out_features=128, bias=True)\n",
      "          (kv): Linear(in_features=128, out_features=256, bias=True)\n",
      "          (cor_embed): Linear(in_features=2, out_features=64, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "          (motion_proj): Linear(in_features=64, out_features=64, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "        (drop_path): Identity()\n",
      "        (norm2): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
      "          (dwconv): DWConv(\n",
      "            (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "          )\n",
      "          (act): GELU(approximate='none')\n",
      "          (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "          (relu): ReLU(inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (norm5): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "    (patch_embed5): OverlapPatchEmbed(\n",
      "      (proj): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "    (block5): ModuleList(\n",
      "      (0-1): 2 x MotionFormerBlock(\n",
      "        (norm1): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "        (attn): InterFrameAttention(\n",
      "          (q): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (kv): Linear(in_features=256, out_features=512, bias=True)\n",
      "          (cor_embed): Linear(in_features=2, out_features=128, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (motion_proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "        (drop_path): Identity()\n",
      "        (norm2): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "          (dwconv): DWConv(\n",
      "            (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "          )\n",
      "          (act): GELU(approximate='none')\n",
      "          (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "          (relu): ReLU(inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (block): ModuleList(\n",
      "    (0): Head(\n",
      "      (upsample): Sequential(\n",
      "        (0): PixelShuffle(upscale_factor=2)\n",
      "        (1): PixelShuffle(upscale_factor=2)\n",
      "      )\n",
      "      (conv): Sequential(\n",
      "        (0): Sequential(\n",
      "          (0): Conv2d(70, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (1): Sequential(\n",
      "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (2): Sequential(\n",
      "          (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=5)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (1): Head(\n",
      "      (upsample): Sequential(\n",
      "        (0): PixelShuffle(upscale_factor=2)\n",
      "        (1): PixelShuffle(upscale_factor=2)\n",
      "      )\n",
      "      (conv): Sequential(\n",
      "        (0): Sequential(\n",
      "          (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (1): Sequential(\n",
      "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (2): Sequential(\n",
      "          (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=5)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (unet): Unet(\n",
      "    (down0): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "      )\n",
      "    )\n",
      "    (down1): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=128)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=128)\n",
      "      )\n",
      "    )\n",
      "    (down2): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=256)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=256)\n",
      "      )\n",
      "    )\n",
      "    (down3): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=512)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=512)\n",
      "      )\n",
      "    )\n",
      "    (up0): Sequential(\n",
      "      (0): ConvTranspose2d(1024, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=256)\n",
      "    )\n",
      "    (up1): Sequential(\n",
      "      (0): ConvTranspose2d(512, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=128)\n",
      "    )\n",
      "    (up2): Sequential(\n",
      "      (0): ConvTranspose2d(256, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (up3): Sequential(\n",
      "      (0): ConvTranspose2d(128, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=32)\n",
      "    )\n",
      "    (conv): Conv2d(32, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# TTA = True\n",
    "# I0_ = (torch.tensor(I0.transpose(2, 0, 1)).cuda() / 255.).unsqueeze(0)\n",
    "# I2_ = (torch.tensor(I2.transpose(2, 0, 1)).cuda() / 255.).unsqueeze(0)\n",
    "\n",
    "# padder = InputPadder(I0_.shape, divisor=32)\n",
    "# I0_, I2_ = padder.pad(I0_, I2_)\n",
    "\n",
    "cfg.MODEL_CONFIG['LOGNAME'] = 'ours_t'\n",
    "cfg.MODEL_CONFIG['MODEL_ARCH'] = cfg.init_model_config(\n",
    "    F=32,\n",
    "    depth=[2, 2, 2, 4, 4]\n",
    ")\n",
    "model = Model(-1,use_perceptual_loss= False)\n",
    "model.load_model()\n",
    "model.eval()\n",
    "model.get_device()\n",
    "print(model.net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MultiScaleFlow(\n",
      "  (feature_bone): MotionFormer(\n",
      "    (block1): ConvBlock(\n",
      "      (conv): Sequential(\n",
      "        (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=16)\n",
      "        (2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (3): PReLU(num_parameters=16)\n",
      "      )\n",
      "    )\n",
      "    (patch_embed2): Sequential(\n",
      "      (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=32)\n",
      "    )\n",
      "    (block2): ConvBlock(\n",
      "      (conv): Sequential(\n",
      "        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=32)\n",
      "        (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (3): PReLU(num_parameters=32)\n",
      "      )\n",
      "    )\n",
      "    (patch_embed3): Sequential(\n",
      "      (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (block3): ConvBlock(\n",
      "      (conv): Sequential(\n",
      "        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "        (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (3): PReLU(num_parameters=64)\n",
      "      )\n",
      "    )\n",
      "    (norm4): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "    (patch_embed4): CrossScalePatchEmbed(\n",
      "      (layers): ModuleList(\n",
      "        (0): Conv2d(64, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(1, 1))\n",
      "        (2): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(2, 2), dilation=(2, 2))\n",
      "        (3): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(1, 1))\n",
      "        (4): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(2, 2), dilation=(2, 2))\n",
      "        (5): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(3, 3), dilation=(3, 3))\n",
      "        (6): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(4, 4), dilation=(4, 4))\n",
      "      )\n",
      "      (proj): Conv2d(112, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "      (norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "    (block4): ModuleList(\n",
      "      (0-1): 2 x MotionFormerBlock(\n",
      "        (norm1): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "        (attn): InterFrameAttention(\n",
      "          (q): Linear(in_features=128, out_features=128, bias=True)\n",
      "          (kv): Linear(in_features=128, out_features=256, bias=True)\n",
      "          (cor_embed): Linear(in_features=2, out_features=64, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "          (motion_proj): Linear(in_features=64, out_features=64, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "        (drop_path): Identity()\n",
      "        (norm2): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
      "          (dwconv): DWConv(\n",
      "            (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "          )\n",
      "          (act): GELU(approximate='none')\n",
      "          (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "          (relu): ReLU(inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (norm5): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "    (patch_embed5): OverlapPatchEmbed(\n",
      "      (proj): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "    (block5): ModuleList(\n",
      "      (0-1): 2 x MotionFormerBlock(\n",
      "        (norm1): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "        (attn): InterFrameAttention(\n",
      "          (q): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (kv): Linear(in_features=256, out_features=512, bias=True)\n",
      "          (cor_embed): Linear(in_features=2, out_features=128, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (motion_proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "        (drop_path): Identity()\n",
      "        (norm2): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "          (dwconv): DWConv(\n",
      "            (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "          )\n",
      "          (act): GELU(approximate='none')\n",
      "          (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "          (relu): ReLU(inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (block): ModuleList(\n",
      "    (0): Head(\n",
      "      (upsample): Sequential(\n",
      "        (0): PixelShuffle(upscale_factor=2)\n",
      "        (1): PixelShuffle(upscale_factor=2)\n",
      "      )\n",
      "      (conv): Sequential(\n",
      "        (0): Sequential(\n",
      "          (0): Conv2d(70, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (1): Sequential(\n",
      "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (2): Sequential(\n",
      "          (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=5)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (1): Head(\n",
      "      (upsample): Sequential(\n",
      "        (0): PixelShuffle(upscale_factor=2)\n",
      "        (1): PixelShuffle(upscale_factor=2)\n",
      "      )\n",
      "      (conv): Sequential(\n",
      "        (0): Sequential(\n",
      "          (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (1): Sequential(\n",
      "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (2): Sequential(\n",
      "          (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=5)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (unet): Unet(\n",
      "    (down0): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "      )\n",
      "    )\n",
      "    (down1): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=128)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=128)\n",
      "      )\n",
      "    )\n",
      "    (down2): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=256)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=256)\n",
      "      )\n",
      "    )\n",
      "    (down3): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=512)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=512)\n",
      "      )\n",
      "    )\n",
      "    (up0): Sequential(\n",
      "      (0): ConvTranspose2d(1024, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=256)\n",
      "    )\n",
      "    (up1): Sequential(\n",
      "      (0): ConvTranspose2d(512, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=128)\n",
      "    )\n",
      "    (up2): Sequential(\n",
      "      (0): ConvTranspose2d(256, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (up3): Sequential(\n",
      "      (0): ConvTranspose2d(128, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=32)\n",
      "    )\n",
      "    (conv): Conv2d(32, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  )\n",
      ")\n",
      ": MultiScaleFlow(\n",
      "  (feature_bone): MotionFormer(\n",
      "    (block1): ConvBlock(\n",
      "      (conv): Sequential(\n",
      "        (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=16)\n",
      "        (2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (3): PReLU(num_parameters=16)\n",
      "      )\n",
      "    )\n",
      "    (patch_embed2): Sequential(\n",
      "      (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=32)\n",
      "    )\n",
      "    (block2): ConvBlock(\n",
      "      (conv): Sequential(\n",
      "        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=32)\n",
      "        (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (3): PReLU(num_parameters=32)\n",
      "      )\n",
      "    )\n",
      "    (patch_embed3): Sequential(\n",
      "      (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (block3): ConvBlock(\n",
      "      (conv): Sequential(\n",
      "        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "        (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (3): PReLU(num_parameters=64)\n",
      "      )\n",
      "    )\n",
      "    (norm4): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "    (patch_embed4): CrossScalePatchEmbed(\n",
      "      (layers): ModuleList(\n",
      "        (0): Conv2d(64, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(1, 1))\n",
      "        (2): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(2, 2), dilation=(2, 2))\n",
      "        (3): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(1, 1))\n",
      "        (4): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(2, 2), dilation=(2, 2))\n",
      "        (5): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(3, 3), dilation=(3, 3))\n",
      "        (6): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(4, 4), dilation=(4, 4))\n",
      "      )\n",
      "      (proj): Conv2d(112, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "      (norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "    (block4): ModuleList(\n",
      "      (0-1): 2 x MotionFormerBlock(\n",
      "        (norm1): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "        (attn): InterFrameAttention(\n",
      "          (q): Linear(in_features=128, out_features=128, bias=True)\n",
      "          (kv): Linear(in_features=128, out_features=256, bias=True)\n",
      "          (cor_embed): Linear(in_features=2, out_features=64, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "          (motion_proj): Linear(in_features=64, out_features=64, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "        (drop_path): Identity()\n",
      "        (norm2): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
      "          (dwconv): DWConv(\n",
      "            (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "          )\n",
      "          (act): GELU(approximate='none')\n",
      "          (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "          (relu): ReLU(inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (norm5): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "    (patch_embed5): OverlapPatchEmbed(\n",
      "      (proj): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "    (block5): ModuleList(\n",
      "      (0-1): 2 x MotionFormerBlock(\n",
      "        (norm1): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "        (attn): InterFrameAttention(\n",
      "          (q): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (kv): Linear(in_features=256, out_features=512, bias=True)\n",
      "          (cor_embed): Linear(in_features=2, out_features=128, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (motion_proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "        (drop_path): Identity()\n",
      "        (norm2): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "          (dwconv): DWConv(\n",
      "            (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "          )\n",
      "          (act): GELU(approximate='none')\n",
      "          (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "          (relu): ReLU(inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (block): ModuleList(\n",
      "    (0): Head(\n",
      "      (upsample): Sequential(\n",
      "        (0): PixelShuffle(upscale_factor=2)\n",
      "        (1): PixelShuffle(upscale_factor=2)\n",
      "      )\n",
      "      (conv): Sequential(\n",
      "        (0): Sequential(\n",
      "          (0): Conv2d(70, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (1): Sequential(\n",
      "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (2): Sequential(\n",
      "          (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=5)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (1): Head(\n",
      "      (upsample): Sequential(\n",
      "        (0): PixelShuffle(upscale_factor=2)\n",
      "        (1): PixelShuffle(upscale_factor=2)\n",
      "      )\n",
      "      (conv): Sequential(\n",
      "        (0): Sequential(\n",
      "          (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (1): Sequential(\n",
      "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=64)\n",
      "        )\n",
      "        (2): Sequential(\n",
      "          (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "          (1): PReLU(num_parameters=5)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (unet): Unet(\n",
      "    (down0): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "      )\n",
      "    )\n",
      "    (down1): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=128)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=128)\n",
      "      )\n",
      "    )\n",
      "    (down2): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=256)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=256)\n",
      "      )\n",
      "    )\n",
      "    (down3): Conv2(\n",
      "      (conv1): Sequential(\n",
      "        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=512)\n",
      "      )\n",
      "      (conv2): Sequential(\n",
      "        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=512)\n",
      "      )\n",
      "    )\n",
      "    (up0): Sequential(\n",
      "      (0): ConvTranspose2d(1024, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=256)\n",
      "    )\n",
      "    (up1): Sequential(\n",
      "      (0): ConvTranspose2d(512, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=128)\n",
      "    )\n",
      "    (up2): Sequential(\n",
      "      (0): ConvTranspose2d(256, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (up3): Sequential(\n",
      "      (0): ConvTranspose2d(128, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=32)\n",
      "    )\n",
      "    (conv): Conv2d(32, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  )\n",
      ")\n",
      "feature_bone: MotionFormer(\n",
      "  (block1): ConvBlock(\n",
      "    (conv): Sequential(\n",
      "      (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=16)\n",
      "      (2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (3): PReLU(num_parameters=16)\n",
      "    )\n",
      "  )\n",
      "  (patch_embed2): Sequential(\n",
      "    (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=32)\n",
      "  )\n",
      "  (block2): ConvBlock(\n",
      "    (conv): Sequential(\n",
      "      (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=32)\n",
      "      (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (3): PReLU(num_parameters=32)\n",
      "    )\n",
      "  )\n",
      "  (patch_embed3): Sequential(\n",
      "    (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=64)\n",
      "  )\n",
      "  (block3): ConvBlock(\n",
      "    (conv): Sequential(\n",
      "      (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "      (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (3): PReLU(num_parameters=64)\n",
      "    )\n",
      "  )\n",
      "  (norm4): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "  (patch_embed4): CrossScalePatchEmbed(\n",
      "    (layers): ModuleList(\n",
      "      (0): Conv2d(64, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(1, 1))\n",
      "      (2): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(2, 2), dilation=(2, 2))\n",
      "      (3): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(1, 1))\n",
      "      (4): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(2, 2), dilation=(2, 2))\n",
      "      (5): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(3, 3), dilation=(3, 3))\n",
      "      (6): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(4, 4), dilation=(4, 4))\n",
      "    )\n",
      "    (proj): Conv2d(112, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "    (norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
      "  )\n",
      "  (block4): ModuleList(\n",
      "    (0-1): 2 x MotionFormerBlock(\n",
      "      (norm1): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "      (attn): InterFrameAttention(\n",
      "        (q): Linear(in_features=128, out_features=128, bias=True)\n",
      "        (kv): Linear(in_features=128, out_features=256, bias=True)\n",
      "        (cor_embed): Linear(in_features=2, out_features=64, bias=True)\n",
      "        (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "        (proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "        (motion_proj): Linear(in_features=64, out_features=64, bias=True)\n",
      "        (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "      )\n",
      "      (drop_path): Identity()\n",
      "      (norm2): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "      (mlp): Mlp(\n",
      "        (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
      "        (dwconv): DWConv(\n",
      "          (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "        )\n",
      "        (act): GELU(approximate='none')\n",
      "        (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
      "        (drop): Dropout(p=0.0, inplace=False)\n",
      "        (relu): ReLU(inplace=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (norm5): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "  (patch_embed5): OverlapPatchEmbed(\n",
      "    (proj): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "  )\n",
      "  (block5): ModuleList(\n",
      "    (0-1): 2 x MotionFormerBlock(\n",
      "      (norm1): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "      (attn): InterFrameAttention(\n",
      "        (q): Linear(in_features=256, out_features=256, bias=True)\n",
      "        (kv): Linear(in_features=256, out_features=512, bias=True)\n",
      "        (cor_embed): Linear(in_features=2, out_features=128, bias=True)\n",
      "        (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "        (proj): Linear(in_features=256, out_features=256, bias=True)\n",
      "        (motion_proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "        (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "      )\n",
      "      (drop_path): Identity()\n",
      "      (norm2): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "      (mlp): Mlp(\n",
      "        (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "        (dwconv): DWConv(\n",
      "          (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "        )\n",
      "        (act): GELU(approximate='none')\n",
      "        (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "        (drop): Dropout(p=0.0, inplace=False)\n",
      "        (relu): ReLU(inplace=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      ")\n",
      "feature_bone.block1: ConvBlock(\n",
      "  (conv): Sequential(\n",
      "    (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=16)\n",
      "    (2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (3): PReLU(num_parameters=16)\n",
      "  )\n",
      ")\n",
      "feature_bone.block1.conv: Sequential(\n",
      "  (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=16)\n",
      "  (2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (3): PReLU(num_parameters=16)\n",
      ")\n",
      "feature_bone.block1.conv.0: Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "feature_bone.block1.conv.1: PReLU(num_parameters=16)\n",
      "feature_bone.block1.conv.2: Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "feature_bone.block1.conv.3: PReLU(num_parameters=16)\n",
      "feature_bone.patch_embed2: Sequential(\n",
      "  (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=32)\n",
      ")\n",
      "feature_bone.patch_embed2.0: Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "feature_bone.patch_embed2.1: PReLU(num_parameters=32)\n",
      "feature_bone.block2: ConvBlock(\n",
      "  (conv): Sequential(\n",
      "    (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=32)\n",
      "    (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (3): PReLU(num_parameters=32)\n",
      "  )\n",
      ")\n",
      "feature_bone.block2.conv: Sequential(\n",
      "  (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=32)\n",
      "  (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (3): PReLU(num_parameters=32)\n",
      ")\n",
      "feature_bone.block2.conv.0: Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "feature_bone.block2.conv.1: PReLU(num_parameters=32)\n",
      "feature_bone.block2.conv.2: Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "feature_bone.block2.conv.3: PReLU(num_parameters=32)\n",
      "feature_bone.patch_embed3: Sequential(\n",
      "  (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=64)\n",
      ")\n",
      "feature_bone.patch_embed3.0: Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "feature_bone.patch_embed3.1: PReLU(num_parameters=64)\n",
      "feature_bone.block3: ConvBlock(\n",
      "  (conv): Sequential(\n",
      "    (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=64)\n",
      "    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (3): PReLU(num_parameters=64)\n",
      "  )\n",
      ")\n",
      "feature_bone.block3.conv: Sequential(\n",
      "  (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=64)\n",
      "  (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (3): PReLU(num_parameters=64)\n",
      ")\n",
      "feature_bone.block3.conv.0: Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "feature_bone.block3.conv.1: PReLU(num_parameters=64)\n",
      "feature_bone.block3.conv.2: Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "feature_bone.block3.conv.3: PReLU(num_parameters=64)\n",
      "feature_bone.norm4: LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "feature_bone.patch_embed4: CrossScalePatchEmbed(\n",
      "  (layers): ModuleList(\n",
      "    (0): Conv2d(64, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (1): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(1, 1))\n",
      "    (2): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(2, 2), dilation=(2, 2))\n",
      "    (3): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(1, 1))\n",
      "    (4): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(2, 2), dilation=(2, 2))\n",
      "    (5): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(3, 3), dilation=(3, 3))\n",
      "    (6): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(4, 4), dilation=(4, 4))\n",
      "  )\n",
      "  (proj): Conv2d(112, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "  (norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
      ")\n",
      "feature_bone.patch_embed4.layers: ModuleList(\n",
      "  (0): Conv2d(64, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (1): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(1, 1))\n",
      "  (2): Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(2, 2), dilation=(2, 2))\n",
      "  (3): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(1, 1))\n",
      "  (4): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(2, 2), dilation=(2, 2))\n",
      "  (5): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(3, 3), dilation=(3, 3))\n",
      "  (6): Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(4, 4), dilation=(4, 4))\n",
      ")\n",
      "feature_bone.patch_embed4.layers.0: Conv2d(64, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "feature_bone.patch_embed4.layers.1: Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(1, 1))\n",
      "feature_bone.patch_embed4.layers.2: Conv2d(32, 16, kernel_size=(3, 3), stride=(4, 4), padding=(2, 2), dilation=(2, 2))\n",
      "feature_bone.patch_embed4.layers.3: Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(1, 1))\n",
      "feature_bone.patch_embed4.layers.4: Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(2, 2), dilation=(2, 2))\n",
      "feature_bone.patch_embed4.layers.5: Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(3, 3), dilation=(3, 3))\n",
      "feature_bone.patch_embed4.layers.6: Conv2d(16, 16, kernel_size=(3, 3), stride=(8, 8), padding=(4, 4), dilation=(4, 4))\n",
      "feature_bone.patch_embed4.proj: Conv2d(112, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "feature_bone.patch_embed4.norm: LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
      "feature_bone.block4: ModuleList(\n",
      "  (0-1): 2 x MotionFormerBlock(\n",
      "    (norm1): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "    (attn): InterFrameAttention(\n",
      "      (q): Linear(in_features=128, out_features=128, bias=True)\n",
      "      (kv): Linear(in_features=128, out_features=256, bias=True)\n",
      "      (cor_embed): Linear(in_features=2, out_features=64, bias=True)\n",
      "      (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "      (proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "      (motion_proj): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "    )\n",
      "    (drop_path): Identity()\n",
      "    (norm2): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "    (mlp): Mlp(\n",
      "      (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
      "      (dwconv): DWConv(\n",
      "        (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "      )\n",
      "      (act): GELU(approximate='none')\n",
      "      (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
      "      (drop): Dropout(p=0.0, inplace=False)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "  )\n",
      ")\n",
      "feature_bone.block4.0: MotionFormerBlock(\n",
      "  (norm1): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "  (attn): InterFrameAttention(\n",
      "    (q): Linear(in_features=128, out_features=128, bias=True)\n",
      "    (kv): Linear(in_features=128, out_features=256, bias=True)\n",
      "    (cor_embed): Linear(in_features=2, out_features=64, bias=True)\n",
      "    (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "    (proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "    (motion_proj): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "  )\n",
      "  (drop_path): Identity()\n",
      "  (norm2): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "  (mlp): Mlp(\n",
      "    (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
      "    (dwconv): DWConv(\n",
      "      (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "    )\n",
      "    (act): GELU(approximate='none')\n",
      "    (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
      "    (drop): Dropout(p=0.0, inplace=False)\n",
      "    (relu): ReLU(inplace=True)\n",
      "  )\n",
      ")\n",
      "feature_bone.block4.0.norm1: LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "feature_bone.block4.0.attn: InterFrameAttention(\n",
      "  (q): Linear(in_features=128, out_features=128, bias=True)\n",
      "  (kv): Linear(in_features=128, out_features=256, bias=True)\n",
      "  (cor_embed): Linear(in_features=2, out_features=64, bias=True)\n",
      "  (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "  (proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "  (motion_proj): Linear(in_features=64, out_features=64, bias=True)\n",
      "  (proj_drop): Dropout(p=0.0, inplace=False)\n",
      ")\n",
      "feature_bone.block4.0.attn.q: Linear(in_features=128, out_features=128, bias=True)\n",
      "feature_bone.block4.0.attn.kv: Linear(in_features=128, out_features=256, bias=True)\n",
      "feature_bone.block4.0.attn.cor_embed: Linear(in_features=2, out_features=64, bias=True)\n",
      "feature_bone.block4.0.attn.attn_drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block4.0.attn.proj: Linear(in_features=128, out_features=128, bias=True)\n",
      "feature_bone.block4.0.attn.motion_proj: Linear(in_features=64, out_features=64, bias=True)\n",
      "feature_bone.block4.0.attn.proj_drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block4.0.drop_path: Identity()\n",
      "feature_bone.block4.0.norm2: LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "feature_bone.block4.0.mlp: Mlp(\n",
      "  (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
      "  (dwconv): DWConv(\n",
      "    (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "  )\n",
      "  (act): GELU(approximate='none')\n",
      "  (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
      "  (drop): Dropout(p=0.0, inplace=False)\n",
      "  (relu): ReLU(inplace=True)\n",
      ")\n",
      "feature_bone.block4.0.mlp.fc1: Linear(in_features=128, out_features=512, bias=True)\n",
      "feature_bone.block4.0.mlp.dwconv: DWConv(\n",
      "  (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      ")\n",
      "feature_bone.block4.0.mlp.dwconv.dwconv: Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "feature_bone.block4.0.mlp.act: GELU(approximate='none')\n",
      "feature_bone.block4.0.mlp.fc2: Linear(in_features=512, out_features=128, bias=True)\n",
      "feature_bone.block4.0.mlp.drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block4.0.mlp.relu: ReLU(inplace=True)\n",
      "feature_bone.block4.1: MotionFormerBlock(\n",
      "  (norm1): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "  (attn): InterFrameAttention(\n",
      "    (q): Linear(in_features=128, out_features=128, bias=True)\n",
      "    (kv): Linear(in_features=128, out_features=256, bias=True)\n",
      "    (cor_embed): Linear(in_features=2, out_features=64, bias=True)\n",
      "    (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "    (proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "    (motion_proj): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "  )\n",
      "  (drop_path): Identity()\n",
      "  (norm2): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "  (mlp): Mlp(\n",
      "    (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
      "    (dwconv): DWConv(\n",
      "      (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "    )\n",
      "    (act): GELU(approximate='none')\n",
      "    (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
      "    (drop): Dropout(p=0.0, inplace=False)\n",
      "    (relu): ReLU(inplace=True)\n",
      "  )\n",
      ")\n",
      "feature_bone.block4.1.norm1: LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "feature_bone.block4.1.attn: InterFrameAttention(\n",
      "  (q): Linear(in_features=128, out_features=128, bias=True)\n",
      "  (kv): Linear(in_features=128, out_features=256, bias=True)\n",
      "  (cor_embed): Linear(in_features=2, out_features=64, bias=True)\n",
      "  (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "  (proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "  (motion_proj): Linear(in_features=64, out_features=64, bias=True)\n",
      "  (proj_drop): Dropout(p=0.0, inplace=False)\n",
      ")\n",
      "feature_bone.block4.1.attn.q: Linear(in_features=128, out_features=128, bias=True)\n",
      "feature_bone.block4.1.attn.kv: Linear(in_features=128, out_features=256, bias=True)\n",
      "feature_bone.block4.1.attn.cor_embed: Linear(in_features=2, out_features=64, bias=True)\n",
      "feature_bone.block4.1.attn.attn_drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block4.1.attn.proj: Linear(in_features=128, out_features=128, bias=True)\n",
      "feature_bone.block4.1.attn.motion_proj: Linear(in_features=64, out_features=64, bias=True)\n",
      "feature_bone.block4.1.attn.proj_drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block4.1.drop_path: Identity()\n",
      "feature_bone.block4.1.norm2: LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
      "feature_bone.block4.1.mlp: Mlp(\n",
      "  (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
      "  (dwconv): DWConv(\n",
      "    (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "  )\n",
      "  (act): GELU(approximate='none')\n",
      "  (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
      "  (drop): Dropout(p=0.0, inplace=False)\n",
      "  (relu): ReLU(inplace=True)\n",
      ")\n",
      "feature_bone.block4.1.mlp.fc1: Linear(in_features=128, out_features=512, bias=True)\n",
      "feature_bone.block4.1.mlp.dwconv: DWConv(\n",
      "  (dwconv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      ")\n",
      "feature_bone.block4.1.mlp.dwconv.dwconv: Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)\n",
      "feature_bone.block4.1.mlp.act: GELU(approximate='none')\n",
      "feature_bone.block4.1.mlp.fc2: Linear(in_features=512, out_features=128, bias=True)\n",
      "feature_bone.block4.1.mlp.drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block4.1.mlp.relu: ReLU(inplace=True)\n",
      "feature_bone.norm5: LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "feature_bone.patch_embed5: OverlapPatchEmbed(\n",
      "  (proj): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      ")\n",
      "feature_bone.patch_embed5.proj: Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "feature_bone.patch_embed5.norm: LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "feature_bone.block5: ModuleList(\n",
      "  (0-1): 2 x MotionFormerBlock(\n",
      "    (norm1): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "    (attn): InterFrameAttention(\n",
      "      (q): Linear(in_features=256, out_features=256, bias=True)\n",
      "      (kv): Linear(in_features=256, out_features=512, bias=True)\n",
      "      (cor_embed): Linear(in_features=2, out_features=128, bias=True)\n",
      "      (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "      (proj): Linear(in_features=256, out_features=256, bias=True)\n",
      "      (motion_proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "      (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "    )\n",
      "    (drop_path): Identity()\n",
      "    (norm2): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "    (mlp): Mlp(\n",
      "      (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "      (dwconv): DWConv(\n",
      "        (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "      )\n",
      "      (act): GELU(approximate='none')\n",
      "      (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "      (drop): Dropout(p=0.0, inplace=False)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "  )\n",
      ")\n",
      "feature_bone.block5.0: MotionFormerBlock(\n",
      "  (norm1): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "  (attn): InterFrameAttention(\n",
      "    (q): Linear(in_features=256, out_features=256, bias=True)\n",
      "    (kv): Linear(in_features=256, out_features=512, bias=True)\n",
      "    (cor_embed): Linear(in_features=2, out_features=128, bias=True)\n",
      "    (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "    (proj): Linear(in_features=256, out_features=256, bias=True)\n",
      "    (motion_proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "    (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "  )\n",
      "  (drop_path): Identity()\n",
      "  (norm2): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "  (mlp): Mlp(\n",
      "    (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "    (dwconv): DWConv(\n",
      "      (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "    )\n",
      "    (act): GELU(approximate='none')\n",
      "    (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "    (drop): Dropout(p=0.0, inplace=False)\n",
      "    (relu): ReLU(inplace=True)\n",
      "  )\n",
      ")\n",
      "feature_bone.block5.0.norm1: LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "feature_bone.block5.0.attn: InterFrameAttention(\n",
      "  (q): Linear(in_features=256, out_features=256, bias=True)\n",
      "  (kv): Linear(in_features=256, out_features=512, bias=True)\n",
      "  (cor_embed): Linear(in_features=2, out_features=128, bias=True)\n",
      "  (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
      "  (motion_proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "  (proj_drop): Dropout(p=0.0, inplace=False)\n",
      ")\n",
      "feature_bone.block5.0.attn.q: Linear(in_features=256, out_features=256, bias=True)\n",
      "feature_bone.block5.0.attn.kv: Linear(in_features=256, out_features=512, bias=True)\n",
      "feature_bone.block5.0.attn.cor_embed: Linear(in_features=2, out_features=128, bias=True)\n",
      "feature_bone.block5.0.attn.attn_drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block5.0.attn.proj: Linear(in_features=256, out_features=256, bias=True)\n",
      "feature_bone.block5.0.attn.motion_proj: Linear(in_features=128, out_features=128, bias=True)\n",
      "feature_bone.block5.0.attn.proj_drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block5.0.drop_path: Identity()\n",
      "feature_bone.block5.0.norm2: LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "feature_bone.block5.0.mlp: Mlp(\n",
      "  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "  (dwconv): DWConv(\n",
      "    (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "  )\n",
      "  (act): GELU(approximate='none')\n",
      "  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "  (drop): Dropout(p=0.0, inplace=False)\n",
      "  (relu): ReLU(inplace=True)\n",
      ")\n",
      "feature_bone.block5.0.mlp.fc1: Linear(in_features=256, out_features=1024, bias=True)\n",
      "feature_bone.block5.0.mlp.dwconv: DWConv(\n",
      "  (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      ")\n",
      "feature_bone.block5.0.mlp.dwconv.dwconv: Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "feature_bone.block5.0.mlp.act: GELU(approximate='none')\n",
      "feature_bone.block5.0.mlp.fc2: Linear(in_features=1024, out_features=256, bias=True)\n",
      "feature_bone.block5.0.mlp.drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block5.0.mlp.relu: ReLU(inplace=True)\n",
      "feature_bone.block5.1: MotionFormerBlock(\n",
      "  (norm1): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "  (attn): InterFrameAttention(\n",
      "    (q): Linear(in_features=256, out_features=256, bias=True)\n",
      "    (kv): Linear(in_features=256, out_features=512, bias=True)\n",
      "    (cor_embed): Linear(in_features=2, out_features=128, bias=True)\n",
      "    (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "    (proj): Linear(in_features=256, out_features=256, bias=True)\n",
      "    (motion_proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "    (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "  )\n",
      "  (drop_path): Identity()\n",
      "  (norm2): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "  (mlp): Mlp(\n",
      "    (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "    (dwconv): DWConv(\n",
      "      (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "    )\n",
      "    (act): GELU(approximate='none')\n",
      "    (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "    (drop): Dropout(p=0.0, inplace=False)\n",
      "    (relu): ReLU(inplace=True)\n",
      "  )\n",
      ")\n",
      "feature_bone.block5.1.norm1: LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "feature_bone.block5.1.attn: InterFrameAttention(\n",
      "  (q): Linear(in_features=256, out_features=256, bias=True)\n",
      "  (kv): Linear(in_features=256, out_features=512, bias=True)\n",
      "  (cor_embed): Linear(in_features=2, out_features=128, bias=True)\n",
      "  (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "  (proj): Linear(in_features=256, out_features=256, bias=True)\n",
      "  (motion_proj): Linear(in_features=128, out_features=128, bias=True)\n",
      "  (proj_drop): Dropout(p=0.0, inplace=False)\n",
      ")\n",
      "feature_bone.block5.1.attn.q: Linear(in_features=256, out_features=256, bias=True)\n",
      "feature_bone.block5.1.attn.kv: Linear(in_features=256, out_features=512, bias=True)\n",
      "feature_bone.block5.1.attn.cor_embed: Linear(in_features=2, out_features=128, bias=True)\n",
      "feature_bone.block5.1.attn.attn_drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block5.1.attn.proj: Linear(in_features=256, out_features=256, bias=True)\n",
      "feature_bone.block5.1.attn.motion_proj: Linear(in_features=128, out_features=128, bias=True)\n",
      "feature_bone.block5.1.attn.proj_drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block5.1.drop_path: Identity()\n",
      "feature_bone.block5.1.norm2: LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
      "feature_bone.block5.1.mlp: Mlp(\n",
      "  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "  (dwconv): DWConv(\n",
      "    (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "  )\n",
      "  (act): GELU(approximate='none')\n",
      "  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "  (drop): Dropout(p=0.0, inplace=False)\n",
      "  (relu): ReLU(inplace=True)\n",
      ")\n",
      "feature_bone.block5.1.mlp.fc1: Linear(in_features=256, out_features=1024, bias=True)\n",
      "feature_bone.block5.1.mlp.dwconv: DWConv(\n",
      "  (dwconv): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      ")\n",
      "feature_bone.block5.1.mlp.dwconv.dwconv: Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1024)\n",
      "feature_bone.block5.1.mlp.act: GELU(approximate='none')\n",
      "feature_bone.block5.1.mlp.fc2: Linear(in_features=1024, out_features=256, bias=True)\n",
      "feature_bone.block5.1.mlp.drop: Dropout(p=0.0, inplace=False)\n",
      "feature_bone.block5.1.mlp.relu: ReLU(inplace=True)\n",
      "block: ModuleList(\n",
      "  (0): Head(\n",
      "    (upsample): Sequential(\n",
      "      (0): PixelShuffle(upscale_factor=2)\n",
      "      (1): PixelShuffle(upscale_factor=2)\n",
      "    )\n",
      "    (conv): Sequential(\n",
      "      (0): Sequential(\n",
      "        (0): Conv2d(70, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "      )\n",
      "      (1): Sequential(\n",
      "        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "      )\n",
      "      (2): Sequential(\n",
      "        (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=5)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (1): Head(\n",
      "    (upsample): Sequential(\n",
      "      (0): PixelShuffle(upscale_factor=2)\n",
      "      (1): PixelShuffle(upscale_factor=2)\n",
      "    )\n",
      "    (conv): Sequential(\n",
      "      (0): Sequential(\n",
      "        (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "      )\n",
      "      (1): Sequential(\n",
      "        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=64)\n",
      "      )\n",
      "      (2): Sequential(\n",
      "        (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): PReLU(num_parameters=5)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      ")\n",
      "block.0: Head(\n",
      "  (upsample): Sequential(\n",
      "    (0): PixelShuffle(upscale_factor=2)\n",
      "    (1): PixelShuffle(upscale_factor=2)\n",
      "  )\n",
      "  (conv): Sequential(\n",
      "    (0): Sequential(\n",
      "      (0): Conv2d(70, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (1): Sequential(\n",
      "      (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (2): Sequential(\n",
      "      (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=5)\n",
      "    )\n",
      "  )\n",
      ")\n",
      "block.0.upsample: Sequential(\n",
      "  (0): PixelShuffle(upscale_factor=2)\n",
      "  (1): PixelShuffle(upscale_factor=2)\n",
      ")\n",
      "block.0.upsample.0: PixelShuffle(upscale_factor=2)\n",
      "block.0.upsample.1: PixelShuffle(upscale_factor=2)\n",
      "block.0.conv: Sequential(\n",
      "  (0): Sequential(\n",
      "    (0): Conv2d(70, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=64)\n",
      "  )\n",
      "  (1): Sequential(\n",
      "    (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=64)\n",
      "  )\n",
      "  (2): Sequential(\n",
      "    (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=5)\n",
      "  )\n",
      ")\n",
      "block.0.conv.0: Sequential(\n",
      "  (0): Conv2d(70, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=64)\n",
      ")\n",
      "block.0.conv.0.0: Conv2d(70, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "block.0.conv.0.1: PReLU(num_parameters=64)\n",
      "block.0.conv.1: Sequential(\n",
      "  (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=64)\n",
      ")\n",
      "block.0.conv.1.0: Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "block.0.conv.1.1: PReLU(num_parameters=64)\n",
      "block.0.conv.2: Sequential(\n",
      "  (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=5)\n",
      ")\n",
      "block.0.conv.2.0: Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "block.0.conv.2.1: PReLU(num_parameters=5)\n",
      "block.1: Head(\n",
      "  (upsample): Sequential(\n",
      "    (0): PixelShuffle(upscale_factor=2)\n",
      "    (1): PixelShuffle(upscale_factor=2)\n",
      "  )\n",
      "  (conv): Sequential(\n",
      "    (0): Sequential(\n",
      "      (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (1): Sequential(\n",
      "      (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (2): Sequential(\n",
      "      (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=5)\n",
      "    )\n",
      "  )\n",
      ")\n",
      "block.1.upsample: Sequential(\n",
      "  (0): PixelShuffle(upscale_factor=2)\n",
      "  (1): PixelShuffle(upscale_factor=2)\n",
      ")\n",
      "block.1.upsample.0: PixelShuffle(upscale_factor=2)\n",
      "block.1.upsample.1: PixelShuffle(upscale_factor=2)\n",
      "block.1.conv: Sequential(\n",
      "  (0): Sequential(\n",
      "    (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=64)\n",
      "  )\n",
      "  (1): Sequential(\n",
      "    (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=64)\n",
      "  )\n",
      "  (2): Sequential(\n",
      "    (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=5)\n",
      "  )\n",
      ")\n",
      "block.1.conv.0: Sequential(\n",
      "  (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=64)\n",
      ")\n",
      "block.1.conv.0.0: Conv2d(49, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "block.1.conv.0.1: PReLU(num_parameters=64)\n",
      "block.1.conv.1: Sequential(\n",
      "  (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=64)\n",
      ")\n",
      "block.1.conv.1.0: Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "block.1.conv.1.1: PReLU(num_parameters=64)\n",
      "block.1.conv.2: Sequential(\n",
      "  (0): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=5)\n",
      ")\n",
      "block.1.conv.2.0: Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "block.1.conv.2.1: PReLU(num_parameters=5)\n",
      "unet: Unet(\n",
      "  (down0): Conv2(\n",
      "    (conv1): Sequential(\n",
      "      (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "    (conv2): Sequential(\n",
      "      (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=64)\n",
      "    )\n",
      "  )\n",
      "  (down1): Conv2(\n",
      "    (conv1): Sequential(\n",
      "      (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=128)\n",
      "    )\n",
      "    (conv2): Sequential(\n",
      "      (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=128)\n",
      "    )\n",
      "  )\n",
      "  (down2): Conv2(\n",
      "    (conv1): Sequential(\n",
      "      (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=256)\n",
      "    )\n",
      "    (conv2): Sequential(\n",
      "      (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=256)\n",
      "    )\n",
      "  )\n",
      "  (down3): Conv2(\n",
      "    (conv1): Sequential(\n",
      "      (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=512)\n",
      "    )\n",
      "    (conv2): Sequential(\n",
      "      (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): PReLU(num_parameters=512)\n",
      "    )\n",
      "  )\n",
      "  (up0): Sequential(\n",
      "    (0): ConvTranspose2d(1024, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=256)\n",
      "  )\n",
      "  (up1): Sequential(\n",
      "    (0): ConvTranspose2d(512, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=128)\n",
      "  )\n",
      "  (up2): Sequential(\n",
      "    (0): ConvTranspose2d(256, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=64)\n",
      "  )\n",
      "  (up3): Sequential(\n",
      "    (0): ConvTranspose2d(128, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=32)\n",
      "  )\n",
      "  (conv): Conv2d(32, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      ")\n",
      "unet.down0: Conv2(\n",
      "  (conv1): Sequential(\n",
      "    (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=64)\n",
      "  )\n",
      "  (conv2): Sequential(\n",
      "    (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=64)\n",
      "  )\n",
      ")\n",
      "unet.down0.conv1: Sequential(\n",
      "  (0): Conv2d(49, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=64)\n",
      ")\n",
      "unet.down0.conv1.0: Conv2d(49, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "unet.down0.conv1.1: PReLU(num_parameters=64)\n",
      "unet.down0.conv2: Sequential(\n",
      "  (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=64)\n",
      ")\n",
      "unet.down0.conv2.0: Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "unet.down0.conv2.1: PReLU(num_parameters=64)\n",
      "unet.down1: Conv2(\n",
      "  (conv1): Sequential(\n",
      "    (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=128)\n",
      "  )\n",
      "  (conv2): Sequential(\n",
      "    (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=128)\n",
      "  )\n",
      ")\n",
      "unet.down1.conv1: Sequential(\n",
      "  (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=128)\n",
      ")\n",
      "unet.down1.conv1.0: Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "unet.down1.conv1.1: PReLU(num_parameters=128)\n",
      "unet.down1.conv2: Sequential(\n",
      "  (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=128)\n",
      ")\n",
      "unet.down1.conv2.0: Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "unet.down1.conv2.1: PReLU(num_parameters=128)\n",
      "unet.down2: Conv2(\n",
      "  (conv1): Sequential(\n",
      "    (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=256)\n",
      "  )\n",
      "  (conv2): Sequential(\n",
      "    (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=256)\n",
      "  )\n",
      ")\n",
      "unet.down2.conv1: Sequential(\n",
      "  (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=256)\n",
      ")\n",
      "unet.down2.conv1.0: Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "unet.down2.conv1.1: PReLU(num_parameters=256)\n",
      "unet.down2.conv2: Sequential(\n",
      "  (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=256)\n",
      ")\n",
      "unet.down2.conv2.0: Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "unet.down2.conv2.1: PReLU(num_parameters=256)\n",
      "unet.down3: Conv2(\n",
      "  (conv1): Sequential(\n",
      "    (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=512)\n",
      "  )\n",
      "  (conv2): Sequential(\n",
      "    (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): PReLU(num_parameters=512)\n",
      "  )\n",
      ")\n",
      "unet.down3.conv1: Sequential(\n",
      "  (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=512)\n",
      ")\n",
      "unet.down3.conv1.0: Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "unet.down3.conv1.1: PReLU(num_parameters=512)\n",
      "unet.down3.conv2: Sequential(\n",
      "  (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=512)\n",
      ")\n",
      "unet.down3.conv2.0: Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "unet.down3.conv2.1: PReLU(num_parameters=512)\n",
      "unet.up0: Sequential(\n",
      "  (0): ConvTranspose2d(1024, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=256)\n",
      ")\n",
      "unet.up0.0: ConvTranspose2d(1024, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "unet.up0.1: PReLU(num_parameters=256)\n",
      "unet.up1: Sequential(\n",
      "  (0): ConvTranspose2d(512, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=128)\n",
      ")\n",
      "unet.up1.0: ConvTranspose2d(512, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "unet.up1.1: PReLU(num_parameters=128)\n",
      "unet.up2: Sequential(\n",
      "  (0): ConvTranspose2d(256, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=64)\n",
      ")\n",
      "unet.up2.0: ConvTranspose2d(256, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "unet.up2.1: PReLU(num_parameters=64)\n",
      "unet.up3: Sequential(\n",
      "  (0): ConvTranspose2d(128, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "  (1): PReLU(num_parameters=32)\n",
      ")\n",
      "unet.up3.0: ConvTranspose2d(128, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "unet.up3.1: PReLU(num_parameters=32)\n",
      "unet.conv: Conv2d(32, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n"
     ]
    }
   ],
   "source": [
    "print(model.net)\n",
    "\n",
    "model_z = model.net\n",
    "for name, module in model_z.named_modules():\n",
    "    print(f\"{name}: {module}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import gdown\n",
    "\n",
    "# # gdown.download_folder(\"https://drive.google.com/drive/folders/1pzl4R0nGajgJJQOIBuSAaIxlfNz70OeJ?usp=sharing\")\n",
    "# gdown.download_folder(\"https://drive.google.com/file/d/10JqEX1C20WkGCkc8uaC9ERDjLvXdzxNi/view?usp=sharing\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0m\u001b[01;34mbaseline_ckpt\u001b[0m/  \u001b[01;34memavfi\u001b[0m/          instantngp_dataloader.py  preprocessing.py\n",
      "\u001b[01;34mdata\u001b[0m/           environment.yml  main.py                   \u001b[01;34m__pycache__\u001b[0m/\n",
      "dataset.zip     image_sort.py    Model_explore.ipynb       readme.md\n"
     ]
    }
   ],
   "source": [
    "ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "exclude_layers = ['attn.kv.']\n",
    "for name,param in model_z.named_parameters():\n",
    "    if any(exclude_layer in name for exclude_layer in exclude_layers):\n",
    "        param.requires_grad = True\n",
    "    else:\n",
    "        param.requires_grad = False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "feature_bone.block4.0.attn.kv.weight\n",
      "feature_bone.block4.0.attn.kv.bias\n",
      "feature_bone.block4.1.attn.kv.weight\n",
      "feature_bone.block4.1.attn.kv.bias\n",
      "feature_bone.block5.0.attn.kv.weight\n",
      "feature_bone.block5.0.attn.kv.bias\n",
      "feature_bone.block5.1.attn.kv.weight\n",
      "feature_bone.block5.1.attn.kv.bias\n"
     ]
    }
   ],
   "source": [
    "for name,param in model_z.named_parameters():\n",
    "     if param.requires_grad == True:\n",
    "        print(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## interpolated: masked psnr."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "idls23",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}