parameter_size.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "03e55a41",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch,os,sys\n",
    "from torch import nn\n",
    "import hat.archs.DRCT_arch as DRCT\n",
    "\n",
    "model = DRCT.DRCT(img_size=64,\n",
    "                 patch_size=1,\n",
    "                 in_chans=3,\n",
    "                 embed_dim=180,\n",
    "                 depths=(6, 6, 6, 6,6,6,6, 6, 6, 6,6,6),\n",
    "                 num_heads=(6, 6, 6, 6,6,6,6, 6, 6, 6,6,6),\n",
    "                 window_size=16,\n",
    "                 compress_ratio=3,\n",
    "                 squeeze_factor=30,\n",
    "                 conv_scale=0.01,\n",
    "                 overlap_ratio=0.5,\n",
    "                 mlp_ratio=2.,\n",
    "                 qkv_bias=True,\n",
    "                 qk_scale=None,\n",
    "                 drop_rate=0.,\n",
    "                 attn_drop_rate=0.,\n",
    "                 drop_path_rate=0.1,\n",
    "                 norm_layer=nn.LayerNorm,\n",
    "                 ape=False,\n",
    "                 patch_norm=True,\n",
    "                 use_checkpoint=False,\n",
    "                 upscale=4,\n",
    "                 img_range=1.,\n",
    "                 upsampler='pixelshuffle',\n",
    "                 resi_connection='1conv',\n",
    "                 gc = 32)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "127482d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:106: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)\n",
      "/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:182: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n",
      "/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:122: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)\n",
      "/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:195: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)\n",
      "Unsupported operator aten::sub encountered 1 time(s)\n",
      "Unsupported operator aten::mul encountered 117 time(s)\n",
      "Unsupported operator aten::div encountered 213 time(s)\n",
      "Unsupported operator aten::add encountered 160 time(s)\n",
      "Unsupported operator aten::softmax encountered 30 time(s)\n",
      "Unsupported operator aten::gelu encountered 30 time(s)\n",
      "Unsupported operator aten::leaky_relu_ encountered 25 time(s)\n",
      "Unsupported operator aten::rand encountered 50 time(s)\n",
      "Unsupported operator aten::floor_ encountered 50 time(s)\n",
      "Unsupported operator aten::pixel_shuffle encountered 2 time(s)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| name                             | #elements or shape   |\n",
      "|:---------------------------------|:---------------------|\n",
      "| model                            | 14.1M                |\n",
      "|  conv_first                      |  5.0K                |\n",
      "|   conv_first.weight              |   (180, 3, 3, 3)     |\n",
      "|   conv_first.bias                |   (180,)             |\n",
      "|  patch_embed                     |  0.4K                |\n",
      "|   patch_embed.norm               |   0.4K               |\n",
      "|    patch_embed.norm.weight       |    (180,)            |\n",
      "|    patch_embed.norm.bias         |    (180,)            |\n",
      "|  layers                          |  13.4M               |\n",
      "|   layers.0                       |   2.2M               |\n",
      "|    layers.0.swin1                |    0.3M              |\n",
      "|    layers.0.adjust1              |    5.8K              |\n",
      "|    layers.0.swin2                |    0.4M              |\n",
      "|    layers.0.adjust2              |    6.8K              |\n",
      "|    layers.0.swin3                |    0.5M              |\n",
      "|    layers.0.adjust3              |    7.8K              |\n",
      "|    layers.0.swin4                |    0.5M              |\n",
      "|    layers.0.adjust4              |    8.9K              |\n",
      "|    layers.0.swin5                |    0.6M              |\n",
      "|    layers.0.adjust5              |    55.6K             |\n",
      "|   layers.1                       |   2.2M               |\n",
      "|    layers.1.swin1                |    0.3M              |\n",
      "|    layers.1.adjust1              |    5.8K              |\n",
      "|    layers.1.swin2                |    0.4M              |\n",
      "|    layers.1.adjust2              |    6.8K              |\n",
      "|    layers.1.swin3                |    0.5M              |\n",
      "|    layers.1.adjust3              |    7.8K              |\n",
      "|    layers.1.swin4                |    0.5M              |\n",
      "|    layers.1.adjust4              |    8.9K              |\n",
      "|    layers.1.swin5                |    0.6M              |\n",
      "|    layers.1.adjust5              |    55.6K             |\n",
      "|   layers.2                       |   2.2M               |\n",
      "|    layers.2.swin1                |    0.3M              |\n",
      "|    layers.2.adjust1              |    5.8K              |\n",
      "|    layers.2.swin2                |    0.4M              |\n",
      "|    layers.2.adjust2              |    6.8K              |\n",
      "|    layers.2.swin3                |    0.5M              |\n",
      "|    layers.2.adjust3              |    7.8K              |\n",
      "|    layers.2.swin4                |    0.5M              |\n",
      "|    layers.2.adjust4              |    8.9K              |\n",
      "|    layers.2.swin5                |    0.6M              |\n",
      "|    layers.2.adjust5              |    55.6K             |\n",
      "|   layers.3                       |   2.2M               |\n",
      "|    layers.3.swin1                |    0.3M              |\n",
      "|    layers.3.adjust1              |    5.8K              |\n",
      "|    layers.3.swin2                |    0.4M              |\n",
      "|    layers.3.adjust2              |    6.8K              |\n",
      "|    layers.3.swin3                |    0.5M              |\n",
      "|    layers.3.adjust3              |    7.8K              |\n",
      "|    layers.3.swin4                |    0.5M              |\n",
      "|    layers.3.adjust4              |    8.9K              |\n",
      "|    layers.3.swin5                |    0.6M              |\n",
      "|    layers.3.adjust5              |    55.6K             |\n",
      "|   layers.4                       |   2.2M               |\n",
      "|    layers.4.swin1                |    0.3M              |\n",
      "|    layers.4.adjust1              |    5.8K              |\n",
      "|    layers.4.swin2                |    0.4M              |\n",
      "|    layers.4.adjust2              |    6.8K              |\n",
      "|    layers.4.swin3                |    0.5M              |\n",
      "|    layers.4.adjust3              |    7.8K              |\n",
      "|    layers.4.swin4                |    0.5M              |\n",
      "|    layers.4.adjust4              |    8.9K              |\n",
      "|    layers.4.swin5                |    0.6M              |\n",
      "|    layers.4.adjust5              |    55.6K             |\n",
      "|   layers.5                       |   2.2M               |\n",
      "|    layers.5.swin1                |    0.3M              |\n",
      "|    layers.5.adjust1              |    5.8K              |\n",
      "|    layers.5.swin2                |    0.4M              |\n",
      "|    layers.5.adjust2              |    6.8K              |\n",
      "|    layers.5.swin3                |    0.5M              |\n",
      "|    layers.5.adjust3              |    7.8K              |\n",
      "|    layers.5.swin4                |    0.5M              |\n",
      "|    layers.5.adjust4              |    8.9K              |\n",
      "|    layers.5.swin5                |    0.6M              |\n",
      "|    layers.5.adjust5              |    55.6K             |\n",
      "|  norm                            |  0.4K                |\n",
      "|   norm.weight                    |   (180,)             |\n",
      "|   norm.bias                      |   (180,)             |\n",
      "|  conv_after_body                 |  0.3M                |\n",
      "|   conv_after_body.weight         |   (180, 180, 3, 3)   |\n",
      "|   conv_after_body.bias           |   (180,)             |\n",
      "|  conv_before_upsample            |  0.1M                |\n",
      "|   conv_before_upsample.0         |   0.1M               |\n",
      "|    conv_before_upsample.0.weight |    (64, 180, 3, 3)   |\n",
      "|    conv_before_upsample.0.bias   |    (64,)             |\n",
      "|  upsample                        |  0.3M                |\n",
      "|   upsample.0                     |   0.1M               |\n",
      "|    upsample.0.weight             |    (256, 64, 3, 3)   |\n",
      "|    upsample.0.bias               |    (256,)            |\n",
      "|   upsample.2                     |   0.1M               |\n",
      "|    upsample.2.weight             |    (256, 64, 3, 3)   |\n",
      "|    upsample.2.bias               |    (256,)            |\n",
      "|  conv_last                       |  1.7K                |\n",
      "|   conv_last.weight               |   (3, 64, 3, 3)      |\n",
      "|   conv_last.bias                 |   (3,)               |\n",
      "Params: 14139579\n",
      "FLOPs: 74642767872\n",
      "Multi-Adds: 37321383936.0\n",
      "Forward Passes: 1\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from fvcore.nn import FlopCountAnalysis, parameter_count_table\n",
    "\n",
    "# 加载模型\n",
    "\n",
    "# 定义输入大小\n",
    "input_tensor = torch.randn(1, 3, 64, 64)\n",
    "\n",
    "# 计算 FLOPs\n",
    "flop_count = FlopCountAnalysis(model, input_tensor)\n",
    "flops = flop_count.total()\n",
    "\n",
    "# 计算参数量\n",
    "params_table = parameter_count_table(model)\n",
    "params = sum(p.numel() for p in model.parameters())\n",
    "\n",
    "# 计算 Multi-Adds 和前向传播次数\n",
    "multi_adds = flops / 2\n",
    "forward_passes = 1\n",
    "\n",
    "# 输出结果\n",
    "print(params_table)\n",
    "print(f\"Params: {params}\")\n",
    "print(f\"FLOPs: {flops}\")\n",
    "print(f\"Multi-Adds: {multi_adds}\")\n",
    "print(f\"Forward Passes: {forward_passes}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6f66af37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA available: True\n",
      "CUDA device count: 5\n",
      "CUDA device name: NVIDIA A100-SXM4-40GB\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "print(\"CUDA available:\", torch.cuda.is_available())\n",
    "print(\"CUDA device count:\", torch.cuda.device_count())\n",
    "print(\"CUDA device name:\", torch.cuda.get_device_name(0))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f00fb985",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:106: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)\n",
      "/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:182: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n",
      "/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:122: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)\n",
      "/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:195: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "CUDA out of memory. Tried to allocate 24.00 MiB (GPU 0; 3.81 GiB total capacity; 2.59 GiB already allocated; 8.44 MiB free; 2.74 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 15\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;66;03m# 计算 FLOPs\u001b[39;00m\n\u001b[1;32m     14\u001b[0m flop_count \u001b[38;5;241m=\u001b[39m FlopCountAnalysis(model, input_tensor)\n\u001b[0;32m---> 15\u001b[0m flops \u001b[38;5;241m=\u001b[39m \u001b[43mflop_count\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtotal\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# 计算参数量\u001b[39;00m\n\u001b[1;32m     18\u001b[0m params_table \u001b[38;5;241m=\u001b[39m parameter_count_table(model)\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/fvcore/nn/jit_analysis.py:248\u001b[0m, in \u001b[0;36mJitModelAnalysis.total\u001b[0;34m(self, module_name)\u001b[0m\n\u001b[1;32m    237\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mtotal\u001b[39m(\u001b[38;5;28mself\u001b[39m, module_name: \u001b[38;5;28mstr\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mint\u001b[39m:\n\u001b[1;32m    238\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    239\u001b[0m \u001b[38;5;124;03m    Returns the total aggregated statistic across all operators\u001b[39;00m\n\u001b[1;32m    240\u001b[0m \u001b[38;5;124;03m    for the requested module.\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    246\u001b[0m \u001b[38;5;124;03m        int : The aggregated statistic.\u001b[39;00m\n\u001b[1;32m    247\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 248\u001b[0m     stats \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_analyze\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    249\u001b[0m     module_name \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcanonical_module_name(module_name)\n\u001b[1;32m    250\u001b[0m     total_count \u001b[38;5;241m=\u001b[39m \u001b[38;5;28msum\u001b[39m(stats\u001b[38;5;241m.\u001b[39mcounts[module_name]\u001b[38;5;241m.\u001b[39mvalues())\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/fvcore/nn/jit_analysis.py:551\u001b[0m, in \u001b[0;36mJitModelAnalysis._analyze\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    549\u001b[0m     \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_warn_trace \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mno_tracer_warning\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m    550\u001b[0m         warnings\u001b[38;5;241m.\u001b[39mfilterwarnings(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mignore\u001b[39m\u001b[38;5;124m\"\u001b[39m, category\u001b[38;5;241m=\u001b[39mTracerWarning)\n\u001b[0;32m--> 551\u001b[0m     graph \u001b[38;5;241m=\u001b[39m \u001b[43m_get_scoped_trace_graph\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_model\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_inputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_aliases\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    553\u001b[0m \u001b[38;5;66;03m# Assures even modules not in the trace graph are initialized to zero count\u001b[39;00m\n\u001b[1;32m    554\u001b[0m counts \u001b[38;5;241m=\u001b[39m {}\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/fvcore/nn/jit_analysis.py:176\u001b[0m, in \u001b[0;36m_get_scoped_trace_graph\u001b[0;34m(module, inputs, aliases)\u001b[0m\n\u001b[1;32m    173\u001b[0m     name \u001b[38;5;241m=\u001b[39m aliases[mod]\n\u001b[1;32m    174\u001b[0m     register_hooks(mod, name)\n\u001b[0;32m--> 176\u001b[0m graph, _ \u001b[38;5;241m=\u001b[39m \u001b[43m_get_trace_graph\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodule\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    178\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m handle \u001b[38;5;129;01min\u001b[39;00m hook_handles:\n\u001b[1;32m    179\u001b[0m     handle\u001b[38;5;241m.\u001b[39mremove()\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/jit/_trace.py:1166\u001b[0m, in \u001b[0;36m_get_trace_graph\u001b[0;34m(f, args, kwargs, strict, _force_outplace, return_inputs, _return_inputs_states)\u001b[0m\n\u001b[1;32m   1164\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(args, \u001b[38;5;28mtuple\u001b[39m):\n\u001b[1;32m   1165\u001b[0m     args \u001b[38;5;241m=\u001b[39m (args,)\n\u001b[0;32m-> 1166\u001b[0m outs \u001b[38;5;241m=\u001b[39m \u001b[43mONNXTracedModule\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstrict\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_force_outplace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreturn_inputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_return_inputs_states\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1167\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m outs\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/jit/_trace.py:127\u001b[0m, in \u001b[0;36mONNXTracedModule.forward\u001b[0;34m(self, *args)\u001b[0m\n\u001b[1;32m    124\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    125\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mtuple\u001b[39m(out_vars)\n\u001b[0;32m--> 127\u001b[0m graph, out \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_C\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_create_graph_by_tracing\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    128\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwrapper\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    129\u001b[0m \u001b[43m    \u001b[49m\u001b[43min_vars\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mmodule_state\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    130\u001b[0m \u001b[43m    \u001b[49m\u001b[43m_create_interpreter_name_lookup_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    131\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstrict\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    132\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_force_outplace\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    133\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    135\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_return_inputs:\n\u001b[1;32m    136\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m graph, outs[\u001b[38;5;241m0\u001b[39m], ret_inputs[\u001b[38;5;241m0\u001b[39m]\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/jit/_trace.py:118\u001b[0m, in \u001b[0;36mONNXTracedModule.forward.<locals>.wrapper\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    116\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_return_inputs_states:\n\u001b[1;32m    117\u001b[0m     inputs_states\u001b[38;5;241m.\u001b[39mappend(_unflatten(in_args, in_desc))\n\u001b[0;32m--> 118\u001b[0m outs\u001b[38;5;241m.\u001b[39mappend(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minner\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtrace_inputs\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m    119\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_return_inputs_states:\n\u001b[1;32m    120\u001b[0m     inputs_states[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m=\u001b[39m (inputs_states[\u001b[38;5;241m0\u001b[39m], trace_inputs)\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/nn/modules/module.py:1128\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1125\u001b[0m     bw_hook \u001b[38;5;241m=\u001b[39m hooks\u001b[38;5;241m.\u001b[39mBackwardHook(\u001b[38;5;28mself\u001b[39m, full_backward_hooks)\n\u001b[1;32m   1126\u001b[0m     \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m bw_hook\u001b[38;5;241m.\u001b[39msetup_input_hook(\u001b[38;5;28minput\u001b[39m)\n\u001b[0;32m-> 1128\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1129\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks:\n\u001b[1;32m   1130\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m hook \u001b[38;5;129;01min\u001b[39;00m (\u001b[38;5;241m*\u001b[39m_global_forward_hooks\u001b[38;5;241m.\u001b[39mvalues(), \u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks\u001b[38;5;241m.\u001b[39mvalues()):\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/nn/modules/module.py:1098\u001b[0m, in \u001b[0;36mModule._slow_forward\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1096\u001b[0m         recording_scopes \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m   1097\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1098\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1099\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m   1100\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m recording_scopes:\n",
      "File \u001b[0;32m/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:846\u001b[0m, in \u001b[0;36mDRCT.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    843\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mupsampler \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mpixelshuffle\u001b[39m\u001b[38;5;124m'\u001b[39m:\n\u001b[1;32m    844\u001b[0m     \u001b[38;5;66;03m# for classical SR\u001b[39;00m\n\u001b[1;32m    845\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv_first(x)\n\u001b[0;32m--> 846\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv_after_body(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward_features\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m) \u001b[38;5;241m+\u001b[39m x\n\u001b[1;32m    847\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv_before_upsample(x)\n\u001b[1;32m    848\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv_last(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mupsample(x))\n",
      "File \u001b[0;32m/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:832\u001b[0m, in \u001b[0;36mDRCT.forward_features\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    829\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpos_drop(x)\n\u001b[1;32m    831\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlayers:\n\u001b[0;32m--> 832\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx_size\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    834\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnorm(x)  \u001b[38;5;66;03m# b seq_len c\u001b[39;00m\n\u001b[1;32m    835\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpatch_unembed(x, x_size)\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/nn/modules/module.py:1128\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1125\u001b[0m     bw_hook \u001b[38;5;241m=\u001b[39m hooks\u001b[38;5;241m.\u001b[39mBackwardHook(\u001b[38;5;28mself\u001b[39m, full_backward_hooks)\n\u001b[1;32m   1126\u001b[0m     \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m bw_hook\u001b[38;5;241m.\u001b[39msetup_input_hook(\u001b[38;5;28minput\u001b[39m)\n\u001b[0;32m-> 1128\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1129\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks:\n\u001b[1;32m   1130\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m hook \u001b[38;5;129;01min\u001b[39;00m (\u001b[38;5;241m*\u001b[39m_global_forward_hooks\u001b[38;5;241m.\u001b[39mvalues(), \u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks\u001b[38;5;241m.\u001b[39mvalues()):\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/nn/modules/module.py:1098\u001b[0m, in \u001b[0;36mModule._slow_forward\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1096\u001b[0m         recording_scopes \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m   1097\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1098\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1099\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m   1100\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m recording_scopes:\n",
      "File \u001b[0;32m/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:299\u001b[0m, in \u001b[0;36mSwinBasedFeatFusionBlock.forward\u001b[0;34m(self, x, xsize)\u001b[0m\n\u001b[1;32m    297\u001b[0m         x2 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpe(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlrelu(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39madjust2(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpue(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mswin2(torch\u001b[38;5;241m.\u001b[39mcat((x, x1), \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m), xsize), xsize))))\n\u001b[1;32m    298\u001b[0m         x3 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpe(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlrelu(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39madjust3(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpue(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mswin3(torch\u001b[38;5;241m.\u001b[39mcat((x, x1, x2), \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m), xsize), xsize))))\n\u001b[0;32m--> 299\u001b[0m         x4 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpe(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlrelu(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39madjust4(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpue(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mswin4\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcat\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx1\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx2\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx3\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mxsize\u001b[49m\u001b[43m)\u001b[49m, xsize))))\n\u001b[1;32m    300\u001b[0m         x5 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpe(           \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39madjust5(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpue(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mswin5(torch\u001b[38;5;241m.\u001b[39mcat((x, x1, x2, x3, x4), \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m), xsize), xsize)))\n\u001b[1;32m    302\u001b[0m \u001b[38;5;66;03m#                 return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))) + x\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/nn/modules/module.py:1128\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1125\u001b[0m     bw_hook \u001b[38;5;241m=\u001b[39m hooks\u001b[38;5;241m.\u001b[39mBackwardHook(\u001b[38;5;28mself\u001b[39m, full_backward_hooks)\n\u001b[1;32m   1126\u001b[0m     \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m bw_hook\u001b[38;5;241m.\u001b[39msetup_input_hook(\u001b[38;5;28minput\u001b[39m)\n\u001b[0;32m-> 1128\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1129\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks:\n\u001b[1;32m   1130\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m hook \u001b[38;5;129;01min\u001b[39;00m (\u001b[38;5;241m*\u001b[39m_global_forward_hooks\u001b[38;5;241m.\u001b[39mvalues(), \u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks\u001b[38;5;241m.\u001b[39mvalues()):\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/nn/modules/module.py:1098\u001b[0m, in \u001b[0;36mModule._slow_forward\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1096\u001b[0m         recording_scopes \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m   1097\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1098\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1099\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m   1100\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m recording_scopes:\n",
      "File \u001b[0;32m/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:401\u001b[0m, in \u001b[0;36mSwinTransformerBlock.forward\u001b[0;34m(self, x, x_size)\u001b[0m\n\u001b[1;32m    399\u001b[0m \u001b[38;5;66;03m# W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size\u001b[39;00m\n\u001b[1;32m    400\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minput_resolution \u001b[38;5;241m==\u001b[39m x_size:\n\u001b[0;32m--> 401\u001b[0m     attn_windows \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mattn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx_windows\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mattn_mask\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# nW*B, window_size*window_size, C\u001b[39;00m\n\u001b[1;32m    402\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    403\u001b[0m     attn_windows \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattn(x_windows, mask\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcalculate_mask(x_size)\u001b[38;5;241m.\u001b[39mto(x\u001b[38;5;241m.\u001b[39mdevice))\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/nn/modules/module.py:1128\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1125\u001b[0m     bw_hook \u001b[38;5;241m=\u001b[39m hooks\u001b[38;5;241m.\u001b[39mBackwardHook(\u001b[38;5;28mself\u001b[39m, full_backward_hooks)\n\u001b[1;32m   1126\u001b[0m     \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m bw_hook\u001b[38;5;241m.\u001b[39msetup_input_hook(\u001b[38;5;28minput\u001b[39m)\n\u001b[0;32m-> 1128\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1129\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks:\n\u001b[1;32m   1130\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m hook \u001b[38;5;129;01min\u001b[39;00m (\u001b[38;5;241m*\u001b[39m_global_forward_hooks\u001b[38;5;241m.\u001b[39mvalues(), \u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks\u001b[38;5;241m.\u001b[39mvalues()):\n",
      "File \u001b[0;32m~/anaconda3/envs/hat/lib/python3.8/site-packages/torch/nn/modules/module.py:1098\u001b[0m, in \u001b[0;36mModule._slow_forward\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1096\u001b[0m         recording_scopes \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m   1097\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1098\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1099\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m   1100\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m recording_scopes:\n",
      "File \u001b[0;32m/raid/ming/SR24/SR24/hat/archs/DRCT_arch.py:186\u001b[0m, in \u001b[0;36mWindowAttention.forward\u001b[0;34m(self, x, mask)\u001b[0m\n\u001b[1;32m    183\u001b[0m q, k, v \u001b[38;5;241m=\u001b[39m qkv[\u001b[38;5;241m0\u001b[39m], qkv[\u001b[38;5;241m1\u001b[39m], qkv[\u001b[38;5;241m2\u001b[39m]  \u001b[38;5;66;03m# make torchscript happy (cannot use tensor as tuple)\u001b[39;00m\n\u001b[1;32m    185\u001b[0m q \u001b[38;5;241m=\u001b[39m q \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mscale\n\u001b[0;32m--> 186\u001b[0m attn \u001b[38;5;241m=\u001b[39m (\u001b[43mq\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m@\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mk\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtranspose\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m)\n\u001b[1;32m    188\u001b[0m relative_position_bias \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrelative_position_bias_table[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrelative_position_index\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)]\u001b[38;5;241m.\u001b[39mview(\n\u001b[1;32m    189\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mwindow_size[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mwindow_size[\u001b[38;5;241m1\u001b[39m], \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mwindow_size[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mwindow_size[\u001b[38;5;241m1\u001b[39m], \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)  \u001b[38;5;66;03m# Wh*Ww,Wh*Ww,nH\u001b[39;00m\n\u001b[1;32m    190\u001b[0m relative_position_bias \u001b[38;5;241m=\u001b[39m relative_position_bias\u001b[38;5;241m.\u001b[39mpermute(\u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m1\u001b[39m)\u001b[38;5;241m.\u001b[39mcontiguous()  \u001b[38;5;66;03m# nH, Wh*Ww, Wh*Ww\u001b[39;00m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 24.00 MiB (GPU 0; 3.81 GiB total capacity; 2.59 GiB already allocated; 8.44 MiB free; 2.74 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision.models as models\n",
    "from fvcore.nn import FlopCountAnalysis, parameter_count_table\n",
    "import os \n",
    "os.environ['CUDA_VISIBLE_DEVICES']='4'\n",
    "# 加载模型\n",
    "# 定义输入大小\n",
    "input_tensor = torch.randn(1, 3, 64, 64).cuda()  # 将输入张量移动到 GPU\n",
    "\n",
    "# 将模型移动到 GPU\n",
    "model = model.cuda()\n",
    "\n",
    "# 计算 FLOPs\n",
    "flop_count = FlopCountAnalysis(model, input_tensor)\n",
    "flops = flop_count.total()\n",
    "\n",
    "# 计算参数量\n",
    "params_table = parameter_count_table(model)\n",
    "params = sum(p.numel() for p in model.parameters())\n",
    "\n",
    "# 计算 Multi-Adds 和前向传播次数\n",
    "multi_adds = flops / 2\n",
    "forward_passes = 1\n",
    "\n",
    "# 输出结果\n",
    "print(params_table)\n",
    "print(f\"Params: {params}\")\n",
    "print(f\"FLOPs: {flops}\")\n",
    "print(f\"Multi-Adds: {multi_adds}\")\n",
    "print(f\"Forward Passes: {forward_passes}\")\n",
    "\n",
    "# 释放 GPU 内存\n",
    "del input_tensor, model\n",
    "torch.cuda.empty_cache()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}