{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!cp -r ../input/pytorch-segmentation-models-lib/ ./\n!pip config set global.disable-pip-version-check true\n!pip install -q ./pytorch-segmentation-models-lib/timm-0.4.12-py3-none-any.whl\n\n\n\n\n","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:24:49.486617Z","iopub.execute_input":"2022-09-21T07:24:49.487717Z","iopub.status.idle":"2022-09-21T07:25:04.628947Z","shell.execute_reply.started":"2022-09-21T07:24:49.487607Z","shell.execute_reply":"2022-09-21T07:25:04.627744Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cp -r ../input/einops-041-wheel/ ./\n!pip config set global.disable-pip-version-check true\n!pip install -q ../input/einops-041-wheel/einops-0.4.1-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:25:04.631795Z","iopub.execute_input":"2022-09-21T07:25:04.632109Z","iopub.status.idle":"2022-09-21T07:25:15.987429Z","shell.execute_reply.started":"2022-09-21T07:25:04.632077Z","shell.execute_reply":"2022-09-21T07:25:15.986235Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\n\n\nsys.path.insert(0, '../input/coat-5layer')\nsys.path.insert(0, '../input/coat5levelswafold3')\nsys.path.insert(0, '../input/coatsmall')\n\n","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:25:15.989427Z","iopub.execute_input":"2022-09-21T07:25:15.989826Z","iopub.status.idle":"2022-09-21T07:25:15.996897Z","shell.execute_reply.started":"2022-09-21T07:25:15.989778Z","shell.execute_reply":"2022-09-21T07:25:15.995317Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import timm\nimport importlib\nfrom timeit import default_timer as timer\n\nimport torch\nimport torch.cuda.amp as amp\nimport torch.nn.functional as F\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport tifffile \nimport cv2\nimport os\nimport gc\nfrom tqdm.notebook import tqdm\nimport rasterio\nfrom rasterio.windows import Window\n\nfrom fastai.vision.all import *\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A # Augmentations\nimport warnings\nwarnings.filterwarnings(\"ignore\")\nimport glob\nimport copy\nprint('import ok\\n')","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:25:16.000011Z","iopub.execute_input":"2022-09-21T07:25:16.000842Z","iopub.status.idle":"2022-09-21T07:25:21.086567Z","shell.execute_reply.started":"2022-09-21T07:25:16.000806Z","shell.execute_reply":"2022-09-21T07:25:21.084679Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport cv2\n\nclass dotdict(dict):\n\t__setattr__ = dict.__setitem__\n\t__delattr__ = dict.__delitem__\n\t\n\tdef __getattr__(self, name):\n\t\ttry:\n\t\t\treturn self[name]\n\t\texcept KeyError:\n\t\t\traise AttributeError(name)\n\n\n#--- helper ----------\ndef time_to_str(t, mode='min'):\n\tif mode=='min':\n\t\tt  = int(t)/60\n\t\thr = t//60\n\t\tmin = t%60\n\t\treturn '%2d hr %02d min'%(hr,min)\n\t\n\telif mode=='sec':\n\t\tt   = int(t)\n\t\tmin = t//60\n\t\tsec = t%60\n\t\treturn '%2d min %02d sec'%(min,sec)\n\t\n\telse:\n\t\traise NotImplementedError\n\t\ndef image_show(name, image, type='bgr', resize=1):\n\tif type == 'rgb': image = np.ascontiguousarray(image[:,:,::-1])\n\tH,W = image.shape[0:2]\n\t\n\tcv2.namedWindow(name, cv2.WINDOW_GUI_NORMAL)  #WINDOW_NORMAL #WINDOW_GUI_EXPANDED\n\tcv2.imshow(name, image) #.astype(np.uint8))\n\tcv2.resizeWindow(name, round(resize*W), round(resize*H))","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:25:21.087852Z","iopub.execute_input":"2022-09-21T07:25:21.088421Z","iopub.status.idle":"2022-09-21T07:25:21.098578Z","shell.execute_reply.started":"2022-09-21T07:25:21.088389Z","shell.execute_reply":"2022-09-21T07:25:21.097508Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys, os\n\nsys.path.append('../input/hubmap-submit-06/[third_party]')\nsys.path.append('../[third_party]')\n#sys.path.append('/root/share1/kaggle/2022/hubmap-organ-segmentation/code/hubmap-dummy-02/[submit]/[third_party]')\n\n# from my_lib_kv import *\n\nimport tifffile as tiff\nimport json\nimport cv2\nimport pandas as pd\nimport math\nimport numpy as np\n\n\n\n##--------------------------------------------------------------------------------------\norgan_meta = dotdict(\n\tkidney = dotdict(\n\t\tlabel = 1,\n\t\tum    = 0.5000,\n\t\tftu   ='glomeruli',\n\t),\n\tprostate = dotdict(\n\t\tlabel = 2,\n\t\tum    = 6.2630,\n\t\tftu   ='glandular acinus',\n\t),\n\tlargeintestine = dotdict(\n\t\tlabel = 3,\n\t\tum    = 0.2290,\n\t\tftu   ='crypt',\n\t),\n\tspleen = dotdict(\n\t\tlabel = 4,\n\t\tum    = 0.4945,\n\t\tftu   ='white pulp',\n\t),\n\tlung = dotdict(\n\t\tlabel = 5,\n\t\tum    = 0.7562,\n\t\tftu   ='alveolus',\n\t),\n)\n\n\n\norgan_to_label = {k: organ_meta[k].label for k in organ_meta.keys()}\nlabel_to_organ = {v:k for k,v in organ_to_label.items()}\nnum_organ=5\n#['kidney', 'prostate', 'largeintestine', 'spleen', 'lung']\n\n\ndef read_tiff(image_file, mode='rgb'):\n\timage = tiff.imread(image_file)\n\timage = image.squeeze()\n\tif image.shape[0] == 3:\n\t\timage = image.transpose(1, 2, 0)\n\tif mode=='bgr':\n\t\timage = image[:,:,::-1]\n\timage = np.ascontiguousarray(image)\n\treturn image\n\ndef read_json_as_list(json_file):\n\twith open(json_file) as f:\n\t\tj = json.load(f)\n\treturn j\n\n\n# # --- rle ---------------------------------\ndef rle_decode(rle, height, width , fill=255, dtype=np.uint8):\n\ts = rle.split()\n\tstart, length = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n\tstart -= 1\n\tmask = np.zeros(height*width, dtype=dtype)\n\tfor i, l in zip(start, length):\n\t\tmask[i:i+l] = fill\n\tmask = mask.reshape(width,height).T\n\tmask = np.ascontiguousarray(mask)\n\treturn mask\n\n\ndef rle_encode(mask):\n\tm = mask.T.flatten()\n\tm = np.concatenate([[0], m, [0]])\n\trun = np.where(m[1:] != m[:-1])[0] + 1\n\trun[1::2] -= run[::2]\n\trle =  ' '.join(str(r) for r in run)\n\treturn rle\n\n\n#\n# # --draw ------------------------------------------\ndef mask_to_inner_contour(mask):\n\tmask = mask>0.5\n\tpad = np.lib.pad(mask, ((1, 1), (1, 1)), 'reflect')\n\tcontour = mask & (\n\t\t\t(pad[1:-1,1:-1] != pad[:-2,1:-1]) \\\n\t\t\t| (pad[1:-1,1:-1] != pad[2:,1:-1]) \\\n\t\t\t| (pad[1:-1,1:-1] != pad[1:-1,:-2]) \\\n\t\t\t| (pad[1:-1,1:-1] != pad[1:-1,2:])\n\t)\n\treturn contour\n\n\ndef draw_contour_overlay(image, mask, color=(0,0,255), thickness=1):\n\tcontour =  mask_to_inner_contour(mask)\n\tif thickness==1:\n\t\timage[contour] = color\n\telse:\n\t\tr = max(1,thickness//2)\n\t\tfor y,x in np.stack(np.where(contour)).T:\n\t\t\tcv2.circle(image, (x,y), r, color, lineType=cv2.LINE_4 )\n\treturn image\n\ndef result_to_overlay(image, mask, probability=None, **kwargs):\n \n\t\n\tH,W,C= image.shape\n\tif mask is None:\n\t\tmask = np.zeros((H,W),np.float32)\n\tif probability is None:\n\t\tprobability = np.zeros((H,W),np.float32)\n\t\t\n\to1 = np.zeros((H,W,3),np.float32)\n\to1[:,:,2] = mask\n\to1[:,:,1] = probability\n\t\n\to2 = image.copy()\n\to2 = o2*0.5\n\to2[:,:,1] += 0.5*probability\n\to2 = draw_contour_overlay(o2, mask, color=(0,0,1), thickness=max(3,int(7*H/1500)))\n\t\n\t#---\n\to2,image,o1 = [(m*255).astype(np.uint8) for m in [o2,image,o1]]\n\tif kwargs.get('dice_score',-1)>=0:\n\t\tdraw_shadow_text(o2,'dice=%0.5f'%kwargs.get('dice_score'),(20,80),2.5,(255,255,255),5)\n\tif kwargs.get('d',None) is not None:\n\t\td = kwargs.get('d')\n\t\tdraw_shadow_text(o2,d['id'],(20,140),1.5,(255,255,255),3)\n\t\tdraw_shadow_text(o2,d.organ+'(%s)'%(organ_meta[d.organ].ftu),(20,190),1.5,(255,255,255),3)\n\t\tdraw_shadow_text(o2,'%0.1f um'%(d.pixel_size),(20,240),1.5,(255,255,255),3)\n\t\ts100 = int(100/d.pixel_size)\n\t\tsx,sy = W-s100-40,H-80\n\t\tcv2.rectangle(o2,(sx,sy),(sx+s100,sy+s100//2),(0,0,0),-1)\n\t\tdraw_shadow_text(o2,'100um',(sx+8,sy+40),1,(255,255,255),2)\n\t\tpass\n\t\n\t#draw_shadow_text(image,'input',(5,15),0.6,(1,1,1),1)\n\t#draw_shadow_text(im_paste,'predict',(5,15),0.6,(1,1,1),1)\n\n\toverlay = np.hstack([o2,image,o1])\n\treturn overlay\n\n# --lb metric ------------------------------------------\n# https://www.kaggle.com/competitions/hubmap-organ-segmentation/overview/supervised-ml-evaluation\n\ndef compute_dice_score(probability, mask):\n\tN = len(probability)\n\tp = probability.reshape(N,-1)\n\tt = mask.reshape(N,-1)\n\t\n\tp = p>0.5\n\tt = t>0.5\n\tuion = p.sum(-1) + t.sum(-1)\n\toverlap = (p*t).sum(-1)\n\tdice = 2*overlap/(uion+0.0001)\n\treturn dice","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:25:21.100109Z","iopub.execute_input":"2022-09-21T07:25:21.100822Z","iopub.status.idle":"2022-09-21T07:25:21.130530Z","shell.execute_reply.started":"2022-09-21T07:25:21.100786Z","shell.execute_reply":"2022-09-21T07:25:21.129266Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nimport pdb\n\nclass MixUpSample(nn.Module):\n\tdef __init__( self, scale_factor=2):\n\t\tsuper().__init__()\n\t\tself.mixing = nn.Parameter(torch.tensor(0.5))\n\t\tself.scale_factor = scale_factor\n\t\n\tdef forward(self, x):\n\t\tx = self.mixing *F.interpolate(x, scale_factor=self.scale_factor, mode='bilinear', align_corners=False) \\\n\t\t\t+ (1-self.mixing )*F.interpolate(x, scale_factor=self.scale_factor, mode='nearest')\n\t\treturn x\n\n#https://github.com/lhoyer/DAFormer/blob/master/mmseg/models/decode_heads/daformer_head.py\ndef Conv2dBnReLU(in_channel, out_channel, kernel_size=3, padding=1,stride=1, dilation=1):\n\treturn nn.Sequential(\n\t\tnn.Conv2d(in_channel, out_channel, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation, bias=False),\n\t\tnn.BatchNorm2d(out_channel),\n\t\tnn.ReLU(inplace=True),\n\t)\n\nclass ASPP(nn.Module):\n\t\n\tdef __init__(self,\n\t\t\t\t in_channel,\n\t\t\t\t channel,\n\t\t\t\t dilation,\n\t\t\t\t ):\n\t\tsuper(ASPP, self).__init__()\n\t\t\n\t\tself.conv = nn.ModuleList()\n\t\tfor d in dilation:\n\t\t\tself.conv.append(\n\t\t\t\tConv2dBnReLU(\n\t\t\t\t\tin_channel,\n\t\t\t\t\tchannel,\n\t\t\t\t\tkernel_size=1 if d == 1 else 3,\n\t\t\t\t\tdilation=d,\n\t\t\t\t\tpadding=0 if d == 1 else d,\n\t\t\t\t)\n\t\t\t)\n\t\t\n\t\tself.out = Conv2dBnReLU(\n\t\t\tlen(dilation) * channel,\n\t\t\tchannel,\n\t\t\tkernel_size=3,\n\t\t\tpadding=1,\n\t\t\t)\n\t\n\tdef forward(self, x):\n\t\taspp = []\n\t\tfor conv in self.conv:\n\t\t\taspp.append(conv(x))\n\t\taspp = torch.cat(aspp, dim=1)\n\t\tout = self.out(aspp)\n\t\treturn out\n\n#DepthwiseSeparable\nclass DSConv2d(nn.Module):\n\tdef __init__(self,\n\t\t\t\t in_channel,\n\t\t\t\t out_channel,\n\t\t\t\t kernel_size,\n\t\t\t\t stride   = 1,\n\t\t\t\t padding  = 0,\n\t\t\t\t dilation = 1\n\t\t):\n\t\tsuper().__init__()\n\t\t\n\t\tself.depthwise = nn.Sequential(\n\t\t\tnn.Conv2d( in_channel, in_channel, kernel_size, stride=stride, padding=padding, dilation=dilation),\n\t\t\tnn.BatchNorm2d(in_channel),\n\t\t\tnn.ReLU(inplace=True)\n\t\t)\n\t\t\n\t\tself.pointwise = nn.Sequential(\n\t\t\tnn.Conv2d( in_channel, out_channel, kernel_size=1, stride=1, padding=0),\n\t\t\tnn.BatchNorm2d(out_channel),\n\t\t\tnn.ReLU(inplace=True)\n\t\t)\n\t\n\tdef forward(self, x):\n\t\tx = self.depthwise(x)\n\t\tx = self.pointwise(x)\n\t\treturn x\n\nclass DSASPP(nn.Module):\n\t\n\tdef __init__(self,\n\t\t\t\t in_channel,\n\t\t\t\t channel,\n\t\t\t\t dilation,\n\t\t\t\t ):\n\t\tsuper(DSASPP, self).__init__()\n\t\t\n\t\tself.conv = nn.ModuleList()\n\t\tfor d in dilation:\n\t\t\tif d == 1:\n\t\t\t\tself.conv.append(\n\t\t\t\t\tConv2dBnReLU(\n\t\t\t\t\t\tin_channel,\n\t\t\t\t\t\tchannel,\n\t\t\t\t\t\tkernel_size=1 if d == 1 else 3,\n\t\t\t\t\t\tdilation=d,\n\t\t\t\t\t\tpadding=0 if d == 1 else d,\n\t\t\t\t\t)\n\t\t\t\t)\n\t\t\telse:\n\t\t\t\tself.conv.append(\n\t\t\t\t\tDSConv2d(\n\t\t\t\t\t\tin_channel,\n\t\t\t\t\t\tchannel,\n\t\t\t\t\t\tkernel_size=3,\n\t\t\t\t\t\tdilation=d,\n\t\t\t\t\t\tpadding=d,\n\t\t\t\t\t)\n\t\t\t\t)\n\t\t\n\t\tself.out = Conv2dBnReLU(\n\t\t\tlen(dilation) * channel,\n\t\t\tchannel,\n\t\t\tkernel_size=3,\n\t\t\tpadding=1,\n\t\t)\n\t \n\tdef forward(self, x):\n\t\taspp = []\n\t\tfor conv in self.conv:\n\t\t\taspp.append(conv(x))\n\t\taspp = torch.cat(aspp, dim=1)\n\t\tout = self.out(aspp)\n\t\treturn out\n\n\t\n##############################################################################################33\n\nclass DaformerDecoder(nn.Module):\n\tdef __init__(\n\t\t\tself,\n\t\t\tencoder_dim = [32, 64, 160, 256],\n\t\t\tdecoder_dim = 256,\n\t\t\tdilation = [1, 6, 12, 18],\n\t\t\tuse_bn_mlp  = True,\n\t\t\tfuse = 'conv3x3',\n\t):\n\t\tsuper().__init__()\n\t\tself.mlp = nn.ModuleList([\n\t\t\tnn.Sequential(\n\t\t\t\t# Conv2dBnReLU(dim, decoder_dim, 1, padding=0), #follow mmseg to use conv-bn-relu\n\t\t\t\t*(\n\t\t\t\t  ( nn.Conv2d(dim, decoder_dim, 1, padding= 0,  bias=False),\n\t\t\t\t\tnn.BatchNorm2d(decoder_dim),\n\t\t\t\t\tnn.ReLU(inplace=True),\n\t\t\t\t)if use_bn_mlp else\n\t\t\t\t  ( nn.Conv2d(dim, decoder_dim, 1, padding= 0,  bias=True),)\n\t\t\t\t),\n\t\t\t\t\n\t\t\t\tMixUpSample(2**i) if i!=0 else nn.Identity(),\n\t\t\t) for i, dim in enumerate(encoder_dim)])\n\t  \n\t\tif fuse=='conv1x1':\n\t\t\tself.fuse = nn.Sequential(\n\t\t\t\tnn.Conv2d(len(encoder_dim) * decoder_dim, decoder_dim, 1, padding=0, bias=False),\n\t\t\t\tnn.BatchNorm2d(decoder_dim),\n\t\t\t\tnn.ReLU(inplace=True),\n\t\t\t)\n\t\t\n\t\tif fuse=='conv3x3':\n\t\t\tself.fuse = nn.Sequential(\n\t\t\t\tnn.Conv2d(len(encoder_dim) * decoder_dim, decoder_dim, 3, padding=1, bias=False),\n\t\t\t\tnn.BatchNorm2d(decoder_dim),\n\t\t\t\tnn.ReLU(inplace=True),\n\t\t\t)\n\t\t\n\t\tif fuse=='aspp':\n\t\t\tself.fuse = ASPP(\n\t\t\t\tdecoder_dim*len(encoder_dim),\n\t\t\t\tdecoder_dim,\n\t\t\t\tdilation,\n\t\t\t)\n\t\t\t\n\t\tif fuse=='ds-aspp':\n\t\t\tself.fuse = DSASPP(\n\t\t\t\tdecoder_dim*len(encoder_dim),\n\t\t\t\tdecoder_dim,\n\t\t\t\tdilation,\n\t\t\t)\n\t\t\n\t\n\tdef forward(self, feature):\n\t\t\n\t\tout = []\n\t\tfor i,f in enumerate(feature):\n\t\t\tf = self.mlp[i](f)\n\t\t\tout.append(f)\n\t\t\t# pdb.set_trace()\n\t\t\t#print(f.shape)\n\t\tx = self.fuse(torch.cat(out, dim = 1))\n\t\treturn x, out\n\n\nclass daformer_conv3x3 (DaformerDecoder):\n\tdef __init__(self, **kwargs):\n\t\tsuper(daformer_conv3x3, self).__init__(\n\t\t\tfuse = 'conv3x3',\n\t\t\t**kwargs\n\t\t)\nclass daformer_conv1x1 (DaformerDecoder):\n\tdef __init__(self, **kwargs):\n\t\tsuper(daformer_conv1x1, self).__init__(\n\t\t\tfuse = 'conv1x1',\n\t\t\t**kwargs\n\t\t)","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:25:21.132190Z","iopub.execute_input":"2022-09-21T07:25:21.133233Z","iopub.status.idle":"2022-09-21T07:25:21.161971Z","shell.execute_reply.started":"2022-09-21T07:25:21.133193Z","shell.execute_reply":"2022-09-21T07:25:21.160533Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# CoAT-5level\n\n# https://github.com/mlpc-ucsd/CoaT/blob/main/src/models/coat.py\n\n\"\"\"\nCoaT architecture.\n\nModified from timm/models/vision_transformer.py\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD\nfrom timm.models.layers import DropPath, to_2tuple, trunc_normal_\nfrom timm.models.registry import register_model\n\nfrom einops import rearrange\nfrom functools import partial\nfrom torch import nn, einsum\nimport pdb\n\n# __all__ = [\n# \t\"coat_tiny\",\n# \t\"coat_mini\",\n# \t\"coat_small\",\n# \t\"coat_lite_tiny\",\n# \t\"coat_lite_mini\",\n# \t\"coat_lite_small\"\n# ]\n\n\n\ndef init_weight(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        #nn.init.xavier_uniform_(m.weight, gain=1)\n        #nn.init.xavier_normal_(m.weight, gain=1)\n        #nn.init.kaiming_uniform_(m.weight, mode='fan_in', nonlinearity='relu')\n        nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')\n        #nn.init.orthogonal_(m.weight, gain=1)\n        if m.bias is not None:\n            m.bias.data.zero_()\n    elif classname.find('Batch') != -1:\n        m.weight.data.normal_(1,0.02)\n        m.bias.data.zero_()\n    elif classname.find('Linear') != -1:\n        nn.init.orthogonal_(m.weight, gain=1)\n        if m.bias is not None:\n            m.bias.data.zero_()\n    elif classname.find('Embedding') != -1:\n        nn.init.orthogonal_(m.weight, gain=1)\n\n\n#---------------------------------------------\n#https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py#L15\nclass LayerNorm2d(nn.Module):\n\tdef __init__(self, dim, eps=1e-6):\n\t\tsuper().__init__()\n\t\tself.dim = dim\n\t\tself.weight = nn.Parameter(torch.ones(dim))\n\t\tself.bias = nn.Parameter(torch.zeros(dim))\n\t\tself.eps = eps\n\t\n\tdef forward(self, x):\n\t\tbatch_size,C,H,W = x.shape\n\t\t#assert C==self.dim, 'C=%d, self.dim=%d'%(C,self.dim)\n\t\t#print('C=%d, self.dim=%d'%(C,self.dim))\n\t\t\n\t\tu = x.mean(1, keepdim=True)\n\t\ts = (x - u).pow(2).mean(1, keepdim=True)\n\t\tx = (x - u) / torch.sqrt(s + self.eps)\n\t\tx = self.weight[:, None, None] * x + self.bias[:, None, None]\n\t\treturn x\n#---------------------------------------------\n\ndef _cfg_coat(url='', **kwargs):\n\treturn {\n\t\t'url': url,\n\t\t'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,\n\t\t'crop_pct': .9, 'interpolation': 'bicubic',\n\t\t'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,\n\t\t'first_conv': 'patch_embed.proj', 'classifier': 'head',\n\t\t**kwargs\n\t}\n\n\nclass Mlp(nn.Module):\n\t\"\"\" Feed-forward network (FFN, a.k.a. MLP) class. \"\"\"\n\tdef __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n\t\tsuper().__init__()\n\t\tout_features = out_features or in_features\n\t\thidden_features = hidden_features or in_features\n\t\tself.fc1 = nn.Linear(in_features, hidden_features)\n\t\tself.act = act_layer()\n\t\tself.fc2 = nn.Linear(hidden_features, out_features)\n\t\tself.drop = nn.Dropout(drop)\n\t\n\tdef forward(self, x):\n\t\tx = self.fc1(x)\n\t\tx = self.act(x)\n\t\tx = self.drop(x)\n\t\tx = self.fc2(x)\n\t\tx = self.drop(x)\n\t\treturn x\n\n\nclass ConvRelPosEnc(nn.Module):\n\t\"\"\" Convolutional relative position encoding. \"\"\"\n\tdef __init__(self, Ch, h, window):\n\t\t\"\"\"\n\t\tInitialization.\n\t\t\tCh: Channels per head.\n\t\t\th: Number of heads.\n\t\t\twindow: Window size(s) in convolutional relative positional encoding. It can have two forms:\n\t\t\t\t\t1. An integer of window size, which assigns all attention heads with the same window size in ConvRelPosEnc.\n\t\t\t\t\t2. A dict mapping window size to #attention head splits (e.g. {window size 1: #attention head split 1, window size 2: #attention head split 2})\n\t\t\t\t\t   It will apply different window size to the attention head splits.\n\t\t\"\"\"\n\t\tsuper().__init__()\n\t\t\n\t\tif isinstance(window, int):\n\t\t\twindow = {window: h}                                                         # Set the same window size for all attention heads.\n\t\t\tself.window = window\n\t\telif isinstance(window, dict):\n\t\t\tself.window = window\n\t\telse:\n\t\t\traise ValueError()\n\t\t\n\t\tself.conv_list = nn.ModuleList()\n\t\tself.head_splits = []\n\t\tfor cur_window, cur_head_split in window.items():\n\t\t\tdilation = 1                                                                 # Use dilation=1 at default.\n\t\t\tpadding_size = (cur_window + (cur_window - 1) * (dilation - 1)) // 2         # Determine padding size. Ref: https://discuss.pytorch.org/t/how-to-keep-the-shape-of-input-and-output-same-when-dilation-conv/14338\n\t\t\tcur_conv = nn.Conv2d(cur_head_split*Ch, cur_head_split*Ch,\n\t\t\t                     kernel_size=(cur_window, cur_window),\n\t\t\t                     padding=(padding_size, padding_size),\n\t\t\t                     dilation=(dilation, dilation),\n\t\t\t                     groups=cur_head_split*Ch,\n\t\t\t                     )\n\t\t\tself.conv_list.append(cur_conv)\n\t\t\tself.head_splits.append(cur_head_split)\n\t\tself.channel_splits = [x*Ch for x in self.head_splits]\n\t\n\tdef forward(self, q, v, size):\n\t\tB, h, N, Ch = q.shape\n\t\tH, W = size\n\t\tassert N == 1 + H * W\n\t\t\n\t\t# Convolutional relative position encoding.\n\t\tq_img = q[:,:,1:,:]                                                              # Shape: [B, h, H*W, Ch].\n\t\tv_img = v[:,:,1:,:]                                                              # Shape: [B, h, H*W, Ch].\n\t\t\n\t\tv_img = rearrange(v_img, 'B h (H W) Ch -> B (h Ch) H W', H=H, W=W)               # Shape: [B, h, H*W, Ch] -> [B, h*Ch, H, W].\n\t\tv_img_list = torch.split(v_img, self.channel_splits, dim=1)                      # Split according to channels.\n\t\tconv_v_img_list = [conv(x) for conv, x in zip(self.conv_list, v_img_list)]\n\t\tconv_v_img = torch.cat(conv_v_img_list, dim=1)\n\t\tconv_v_img = rearrange(conv_v_img, 'B (h Ch) H W -> B h (H W) Ch', h=h)          # Shape: [B, h*Ch, H, W] -> [B, h, H*W, Ch].\n\t\t\n\t\tEV_hat_img = q_img * conv_v_img\n\t\tzero = torch.zeros((B, h, 1, Ch), dtype=q.dtype, layout=q.layout, device=q.device)\n\t\tEV_hat = torch.cat((zero, EV_hat_img), dim=2)                                # Shape: [B, h, N, Ch].\n\t\t\n\t\treturn EV_hat\n\n\nclass FactorAtt_ConvRelPosEnc(nn.Module):\n\t\"\"\" Factorized attention with convolutional relative position encoding class. \"\"\"\n\tdef __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., shared_crpe=None):\n\t\tsuper().__init__()\n\t\tself.num_heads = num_heads\n\t\thead_dim = dim // num_heads\n\t\tself.scale = qk_scale or head_dim ** -0.5\n\t\t\n\t\tself.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n\t\tself.attn_drop = nn.Dropout(attn_drop)                                           # Note: attn_drop is actually not used.\n\t\tself.proj = nn.Linear(dim, dim)\n\t\tself.proj_drop = nn.Dropout(proj_drop)\n\t\t\n\t\t# Shared convolutional relative position encoding.\n\t\tself.crpe = shared_crpe\n\t\n\tdef forward(self, x, size):\n\t\tB, N, C = x.shape\n\t\t\n\t\t# Generate Q, K, V.\n\t\tqkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)  # Shape: [3, B, h, N, Ch].\n\t\tq, k, v = qkv[0], qkv[1], qkv[2]                                                 # Shape: [B, h, N, Ch].\n\t\t\n\t\t# Factorized attention.\n\t\tk_softmax = k.softmax(dim=2)                                                     # Softmax on dim N.\n\t\tk_softmax_T_dot_v = einsum('b h n k, b h n v -> b h k v', k_softmax, v)          # Shape: [B, h, Ch, Ch].\n\t\tfactor_att        = einsum('b h n k, b h k v -> b h n v', q, k_softmax_T_dot_v)  # Shape: [B, h, N, Ch].\n\t\t\n\t\t# Convolutional relative position encoding.\n\t\tcrpe = self.crpe(q, v, size=size)                                                # Shape: [B, h, N, Ch].\n\t\t\n\t\t# Merge and reshape.\n\t\tx = self.scale * factor_att + crpe\n\t\tx = x.transpose(1, 2).reshape(B, N, C)                                           # Shape: [B, h, N, Ch] -> [B, N, h, Ch] -> [B, N, C].\n\t\t\n\t\t# Output projection.\n\t\tx = self.proj(x)\n\t\tx = self.proj_drop(x)\n\t\t\n\t\treturn x                                                                         # Shape: [B, N, C].\n\n\nclass ConvPosEnc(nn.Module):\n\t\"\"\" Convolutional Position Encoding.\n\t\tNote: This module is similar to the conditional position encoding in CPVT.\n\t\"\"\"\n\tdef __init__(self, dim, k=3):\n\t\tsuper(ConvPosEnc, self).__init__()\n\t\tself.proj = nn.Conv2d(dim, dim, k, 1, k//2, groups=dim)\n\t\n\tdef forward(self, x, size):\n\t\tB, N, C = x.shape\n\t\tH, W = size\n\t\tassert N == 1 + H * W\n\t\t\n\t\t# Extract CLS token and image tokens.\n\t\tcls_token, img_tokens = x[:, :1], x[:, 1:]                                       # Shape: [B, 1, C], [B, H*W, C].\n\t\t\n\t\t# Depthwise convolution.\n\t\tfeat = img_tokens.transpose(1, 2).view(B, C, H, W)\n\t\tx = self.proj(feat) + feat\n\t\tx = x.flatten(2).transpose(1, 2)\n\t\t\n\t\t# Combine with CLS token.\n\t\tx = torch.cat((cls_token, x), dim=1)\n\t\t\n\t\treturn x\n\n\nclass SerialBlock(nn.Module):\n\t\"\"\" Serial block class.\n\t\tNote: In this implementation, each serial block only contains a conv-attention and a FFN (MLP) module. \"\"\"\n\tdef __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,\n\t             drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,\n\t             shared_cpe=None, shared_crpe=None):\n\t\tsuper().__init__()\n\t\t\n\t\t# Conv-Attention.\n\t\tself.cpe = shared_cpe\n\t\t\n\t\tself.norm1 = norm_layer(dim)\n\t\tself.factoratt_crpe = FactorAtt_ConvRelPosEnc(\n\t\t\tdim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n\t\t\tshared_crpe=shared_crpe)\n\t\tself.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n\t\t\n\t\t# MLP.\n\t\tself.norm2 = norm_layer(dim)\n\t\tmlp_hidden_dim = int(dim * mlp_ratio)\n\t\tself.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\t\n\tdef forward(self, x, size):\n\t\t# Conv-Attention.\n\t\tx = self.cpe(x, size)                  # Apply convolutional position encoding.\n\t\tcur = self.norm1(x)\n\t\tcur = self.factoratt_crpe(cur, size)   # Apply factorized attention and convolutional relative position encoding.\n\t\tx = x + self.drop_path(cur)\n\t\t\n\t\t# MLP.\n\t\tcur = self.norm2(x)\n\t\tcur = self.mlp(cur)\n\t\tx = x + self.drop_path(cur)\n\t\t\n\t\treturn x\n\n\nclass ParallelBlock(nn.Module):\n\t\"\"\" Parallel block class. \"\"\"\n\tdef __init__(self, dims, num_heads, mlp_ratios=[], qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,\n\t             drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,\n\t             shared_cpes=None, shared_crpes=None):\n\t\tsuper().__init__()\n\t\t\n\t\t# Conv-Attention.\n\t\tself.cpes = shared_cpes\n\t\t\n\t\tself.norm12 = norm_layer(dims[1])\n\t\tself.norm13 = norm_layer(dims[2])\n\t\tself.norm14 = norm_layer(dims[3])\n\t\tself.norm15 = norm_layer(dims[4])\n\n\n\t\tself.factoratt_crpe2 = FactorAtt_ConvRelPosEnc(\n\t\t\tdims[1], num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n\t\t\tshared_crpe=shared_crpes[1]\n\t\t)\n\t\tself.factoratt_crpe3 = FactorAtt_ConvRelPosEnc(\n\t\t\tdims[2], num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n\t\t\tshared_crpe=shared_crpes[2]\n\t\t)\n\t\tself.factoratt_crpe4 = FactorAtt_ConvRelPosEnc(\n\t\t\tdims[3], num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n\t\t\tshared_crpe=shared_crpes[3]\n\t\t)\n\t\tself.factoratt_crpe5 = FactorAtt_ConvRelPosEnc(\n\t\t\tdims[4], num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n\t\t\tshared_crpe=shared_crpes[4]\n\t\t)\n\n\n\t\tself.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n\t\t\n\t\t# MLP.\n\t\tself.norm22 = norm_layer(dims[1])\n\t\tself.norm23 = norm_layer(dims[2])\n\t\tself.norm24 = norm_layer(dims[3])\n\t\tself.norm25 = norm_layer(dims[4])\n\n\t\tassert dims[1] == dims[2] == dims[3] ==dims[4]                             # In parallel block, we assume dimensions are the same and share the linear transformation.\n\t\tassert mlp_ratios[1] == mlp_ratios[2] == mlp_ratios[3]\n\t\tmlp_hidden_dim = int(dims[1] * mlp_ratios[1])\n\t\tself.mlp2 = self.mlp3 = self.mlp4 =self.mlp5= Mlp(in_features=dims[1], hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\t\n\tdef upsample(self, x, output_size, size):\n\t\t\"\"\" Feature map up-sampling. \"\"\"\n\t\treturn self.interpolate(x, output_size=output_size, size=size)\n\t\n\tdef downsample(self, x, output_size, size):\n\t\t\"\"\" Feature map down-sampling. \"\"\"\n\t\treturn self.interpolate(x, output_size=output_size, size=size)\n\t\n\tdef interpolate(self, x, output_size, size):\n\t\t\"\"\" Feature map interpolation. \"\"\"\n\t\tB, N, C = x.shape\n\t\tH, W = size\n\t\tassert N == 1 + H * W\n\t\t\n\t\tcls_token  = x[:, :1, :]\n\t\timg_tokens = x[:, 1:, :]\n\t\t\n\t\timg_tokens = img_tokens.transpose(1, 2).reshape(B, C, H, W)\n\t\timg_tokens = F.interpolate(img_tokens, size=output_size, mode='bilinear')  # FIXME: May have alignment issue.\n\t\timg_tokens = img_tokens.reshape(B, C, -1).transpose(1, 2)\n\t\t\n\t\tout = torch.cat((cls_token, img_tokens), dim=1)\n\t\t\n\t\treturn out\n\t\n\tdef forward(self, x1, x2, x3, x4, x5,sizes):\n\t\t_, (H2, W2), (H3, W3), (H4, W4),(H5,W5) = sizes\n\t\t\n\t\t# Conv-Attention.\n\t\tx2 = self.cpes[1](x2, size=(H2, W2))  # Note: x1 is ignored.\n\t\tx3 = self.cpes[2](x3, size=(H3, W3))\n\t\tx4 = self.cpes[3](x4, size=(H4, W4))\n\t\tx5 = self.cpes[4](x5, size=(H5, W5))\n\t\t\n\t\tcur2 = self.norm12(x2)\n\t\tcur3 = self.norm13(x3)\n\t\tcur4 = self.norm14(x4)\n\t\tcur5 = self.norm15(x5)\n\n\t\tcur2 = self.factoratt_crpe2(cur2, size=(H2,W2))\n\t\tcur3 = self.factoratt_crpe3(cur3, size=(H3,W3))\n\t\tcur4 = self.factoratt_crpe4(cur4, size=(H4,W4))\n\t\tcur5 = self.factoratt_crpe5(cur5, size=(H5,W5))\n\n\n\t\tupsample3_2 = self.upsample(cur3, output_size=(H2,W2), size=(H3,W3))\n\t\tupsample4_3 = self.upsample(cur4, output_size=(H3,W3), size=(H4,W4))\n\t\tupsample4_2 = self.upsample(cur4, output_size=(H2,W2), size=(H4,W4))\n\t\tdownsample2_3 = self.downsample(cur2, output_size=(H3,W3), size=(H2,W2))\n\t\tdownsample3_4 = self.downsample(cur3, output_size=(H4,W4), size=(H3,W3))\n\t\tdownsample2_4 = self.downsample(cur2, output_size=(H4,W4), size=(H2,W2))\n\t\tupsample5_2 = self.upsample(cur5, output_size=(H2,W2), size=(H5,W5))\n\t\tupsample5_3 = self.upsample(cur5, output_size=(H3,W3), size=(H5,W5))\n\t\tdownsample3_5 = self.downsample(cur3, output_size=(H5,W5), size=(H3,W3))\n\t\tupsample5_4 = self.upsample(cur5, output_size=(H4,W4), size=(H5,W5))\n\t\tdownsample2_5 = self.downsample(cur2, output_size=(H5,W5), size=(H2,W2))\n\t\tdownsample4_5 = self.downsample(cur4, output_size=(H5,W5), size=(H4,W4))\n        \n     \n\n\t\t# cur2 = cur2  + upsample3_2   + upsample4_2+  upsample5_2\n\t\t# cur3 = cur3  + upsample4_3   + downsample2_3 + upsample5_3\n\t\t# cur4 = cur4  + downsample3_4 + downsample2_4 + upsample5_4 \n\t\t# cur5 = cur5  + downsample3_5 + downsample2_5 + downsample4_5 4\n        # not use all connetction\n\t\tcur2 = cur2  + upsample3_2   + upsample4_2\n\t\tcur3 = cur3  + upsample4_3   + downsample2_3\n\t\tcur4 = cur4  + upsample5_4   + downsample2_4\n\t\tcur5 = cur5  + downsample4_5 + downsample2_5\n\n\n\t\tx2 = x2 + self.drop_path(cur2)\n\t\tx3 = x3 + self.drop_path(cur3)\n\t\tx4 = x4 + self.drop_path(cur4)\n\t\tx5 = x5 + self.drop_path(cur5)\n\t\t\n\t\t# MLP.\n\t\tcur2 = self.norm22(x2)\n\t\tcur3 = self.norm23(x3)\n\t\tcur4 = self.norm24(x4)\n\t\tcur5 = self.norm25(x5)\n\n\t\tcur2 = self.mlp2(cur2)\n\t\tcur3 = self.mlp3(cur3)\n\t\tcur4 = self.mlp4(cur4)\n\t\tcur5 = self.mlp5(cur5)\n\n\t\tx2 = x2 + self.drop_path(cur2)\n\t\tx3 = x3 + self.drop_path(cur3)\n\t\tx4 = x4 + self.drop_path(cur4)\n\t\tx5 = x5 + self.drop_path(cur5)\n\t\t\n\t\treturn x1, x2, x3, x4, x5\n\n\nclass PatchEmbed(nn.Module):\n\t\"\"\" Image to Patch Embedding \"\"\"\n\tdef __init__(self, patch_size=16, in_chans=3, embed_dim=768):\n\t\tsuper().__init__()\n\t\tpatch_size = to_2tuple(patch_size)\n\t\t\n\t\tself.patch_size = patch_size\n\t\tself.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)\n\t\tself.norm = nn.LayerNorm(embed_dim)\n\t\n\tdef forward(self, x):\n\t\t_, _, H, W = x.shape\n\t\tout_H, out_W = H // self.patch_size[0], W // self.patch_size[1]\n\t\t\n\t\tx = self.proj(x).flatten(2).transpose(1, 2)\n\t\tout = self.norm(x)\n\t\t\n\t\treturn out, (out_H, out_W)\n\n\nclass CoaT(nn.Module):\n\t\"\"\" CoaT class. \"\"\"\n\tdef __init__(self, patch_size=16, in_chans=3, embed_dims=[0, 0, 0, 0],\n\t             serial_depths=[0, 0, 0, 0], parallel_depth=0,\n\t             num_heads=0, mlp_ratios=[0, 0, 0, 0], qkv_bias=True, qk_scale=None, drop_rate=0., attn_drop_rate=0.,\n\t             drop_path_rate=0.,\n\t             norm_layer=partial(nn.LayerNorm, eps=1e-6),\n\t             return_interm_layers=True,\n\t             out_features=['x1_nocls','x2_nocls','x3_nocls','x4_nocls','x5_nocls'],\n\t             crpe_window={3:2, 5:3, 7:3},\n\t             pretrain=None,\n\t             out_norm = nn.Identity, #use nn.Identity, nn.BatchNorm2d, LayerNorm2d\n\t             **kwargs):\n\t\tsuper().__init__()\n\t\tself.return_interm_layers = return_interm_layers\n\t\tself.pretrain     = pretrain\n\t\tself.embed_dims   = embed_dims\n\t\tself.out_features = out_features\n\t\t#self.num_classes  = num_classes\n\t\t\n\t\t# Patch embeddings.\n\t\tself.patch_embed1 = PatchEmbed(patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dims[0])\n\t\tself.patch_embed2 = PatchEmbed(patch_size=2, in_chans=embed_dims[0], embed_dim=embed_dims[1])\n\t\tself.patch_embed3 = PatchEmbed(patch_size=2, in_chans=embed_dims[1], embed_dim=embed_dims[2])\n\t\tself.patch_embed4 = PatchEmbed(patch_size=2, in_chans=embed_dims[2], embed_dim=embed_dims[3])\n\t\tself.patch_embed5 = PatchEmbed(patch_size=2, in_chans=embed_dims[3], embed_dim=embed_dims[4])\n\t\t\n\t\t# Class tokens.\n\t\tself.cls_token1 = nn.Parameter(torch.zeros(1, 1, embed_dims[0]))\n\t\tself.cls_token2 = nn.Parameter(torch.zeros(1, 1, embed_dims[1]))\n\t\tself.cls_token3 = nn.Parameter(torch.zeros(1, 1, embed_dims[2]))\n\t\tself.cls_token4 = nn.Parameter(torch.zeros(1, 1, embed_dims[3]))\n\t\tself.cls_token5 = nn.Parameter(torch.zeros(1, 1, embed_dims[4]))\n\t\t\n\t\t# Convolutional position encodings.\n\t\tself.cpe1 = ConvPosEnc(dim=embed_dims[0], k=3)\n\t\tself.cpe2 = ConvPosEnc(dim=embed_dims[1], k=3)\n\t\tself.cpe3 = ConvPosEnc(dim=embed_dims[2], k=3)\n\t\tself.cpe4 = ConvPosEnc(dim=embed_dims[3], k=3)\n\t\tself.cpe5 = ConvPosEnc(dim=embed_dims[4], k=3)\n\t\t\n\t\t# Convolutional relative position encodings.\n\t\tself.crpe1 = ConvRelPosEnc(Ch=embed_dims[0] // num_heads, h=num_heads, window=crpe_window)\n\t\tself.crpe2 = ConvRelPosEnc(Ch=embed_dims[1] // num_heads, h=num_heads, window=crpe_window)\n\t\tself.crpe3 = ConvRelPosEnc(Ch=embed_dims[2] // num_heads, h=num_heads, window=crpe_window)\n\t\tself.crpe4 = ConvRelPosEnc(Ch=embed_dims[3] // num_heads, h=num_heads, window=crpe_window)\n\t\tself.crpe5 = ConvRelPosEnc(Ch=embed_dims[4] // num_heads, h=num_heads, window=crpe_window)\n\n\t\t# Enable stochastic depth.\n\t\tdpr = drop_path_rate\n\t\t\n\t\t# Serial blocks 1.\n\t\tself.serial_blocks1 = nn.ModuleList([\n\t\t\tSerialBlock(\n\t\t\t\tdim=embed_dims[0], num_heads=num_heads, mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale,\n\t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n\t\t\t\tshared_cpe=self.cpe1, shared_crpe=self.crpe1\n\t\t\t)\n\t\t\tfor _ in range(serial_depths[0])]\n\t\t)\n\t\t\n\t\t# Serial blocks 2.\n\t\tself.serial_blocks2 = nn.ModuleList([\n\t\t\tSerialBlock(\n\t\t\t\tdim=embed_dims[1], num_heads=num_heads, mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale,\n\t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n\t\t\t\tshared_cpe=self.cpe2, shared_crpe=self.crpe2\n\t\t\t)\n\t\t\tfor _ in range(serial_depths[1])]\n\t\t)\n\t\t\n\t\t# Serial blocks 3.\n\t\tself.serial_blocks3 = nn.ModuleList([\n\t\t\tSerialBlock(\n\t\t\t\tdim=embed_dims[2], num_heads=num_heads, mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale,\n\t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n\t\t\t\tshared_cpe=self.cpe3, shared_crpe=self.crpe3\n\t\t\t)\n\t\t\tfor _ in range(serial_depths[2])]\n\t\t)\n\t\t\n\t\t# Serial blocks 4.\n\t\tself.serial_blocks4 = nn.ModuleList([\n\t\t\tSerialBlock(\n\t\t\t\tdim=embed_dims[3], num_heads=num_heads, mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale,\n\t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n\t\t\t\tshared_cpe=self.cpe4, shared_crpe=self.crpe4\n\t\t\t)\n\t\t\tfor _ in range(serial_depths[3])]\n\t\t)\n\n\t\t# Serial blocks 5.\n\t\tself.serial_blocks5 = nn.ModuleList([\n\t\t\tSerialBlock(\n\t\t\t\tdim=embed_dims[4], num_heads=num_heads, mlp_ratio=mlp_ratios[4], qkv_bias=qkv_bias, qk_scale=qk_scale,\n\t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n\t\t\t\tshared_cpe=self.cpe5, shared_crpe=self.crpe5\n\t\t\t)\n\t\t\tfor _ in range(serial_depths[4])]\n\t\t)\n\t\t\n\t\t# Parallel blocks.\n\t\tself.parallel_depth = parallel_depth\n\t\tif self.parallel_depth > 0:\n\t\t\tself.parallel_blocks = nn.ModuleList([\n\t\t\t\tParallelBlock(\n\t\t\t\t\tdims=embed_dims, num_heads=num_heads, mlp_ratios=mlp_ratios, qkv_bias=qkv_bias, qk_scale=qk_scale,\n\t\t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n\t\t\t\t\tshared_cpes=[self.cpe1, self.cpe2, self.cpe3, self.cpe4,  self.cpe5],\n\t\t\t\t\tshared_crpes=[self.crpe1, self.crpe2, self.crpe3, self.crpe4,  self.crpe5]\n\t\t\t\t)\n\t\t\t\tfor _ in range(parallel_depth)]\n\t\t\t)\n\t\t\n\t\t# Classification head(s).\n\t\t# if not self.return_interm_layers:\n\t\t# \tself.norm1 = norm_layer(embed_dims[0])\n\t\t# \tself.norm2 = norm_layer(embed_dims[1])\n\t\t# \tself.norm3 = norm_layer(embed_dims[2])\n\t\t# \tself.norm4 = norm_layer(embed_dims[3])\n\t\t#\n\t\t# \tif self.parallel_depth > 0:                                  # CoaT series: Aggregate features of last three scales for classification.\n\t\t# \t\tassert embed_dims[1] == embed_dims[2] == embed_dims[3]\n\t\t# \t\tself.aggregate = torch.nn.Conv1d(in_channels=3, out_channels=1, kernel_size=1)\n\t\t# \t\tself.head = nn.Linear(embed_dims[3], num_classes)\n\t\t# \telse:\n\t\t# \t\tself.head = nn.Linear(embed_dims[3], num_classes)        # CoaT-Lite series: Use feature of last scale for classification.\n\t\t#---\n\t\t# add a norm layer for each output\n\t\tself.out_norm = nn.ModuleList(\n\t\t\t[ out_norm(embed_dims[i]) for i in range(4)]\n\t\t)\n\t\t\n\t\t# Initialize weights.\n\t\ttrunc_normal_(self.cls_token1, std=.02)\n\t\ttrunc_normal_(self.cls_token2, std=.02)\n\t\ttrunc_normal_(self.cls_token3, std=.02)\n\t\ttrunc_normal_(self.cls_token4, std=.02)\n\t\ttrunc_normal_(self.cls_token5, std=.02)\n\t\tself.apply(self._init_weights)\n\t\n\tdef _init_weights(self, m):\n\t\tif isinstance(m, nn.Linear):\n\t\t\ttrunc_normal_(m.weight, std=.02)\n\t\t\tif isinstance(m, nn.Linear) and m.bias is not None:\n\t\t\t\tnn.init.constant_(m.bias, 0)\n\t\telif isinstance(m, nn.LayerNorm):\n\t\t\tnn.init.constant_(m.bias, 0)\n\t\t\tnn.init.constant_(m.weight, 1.0)\n\t\n\t@torch.jit.ignore\n\tdef no_weight_decay(self):\n\t\treturn {'cls_token1', 'cls_token2', 'cls_token3', 'cls_token4'}\n\t\n\t# def get_classifier(self):\n\t# \treturn self.head\n\t\n\t# def reset_classifier(self, num_classes, global_pool=''):\n\t# \tself.num_classes = num_classes\n\t# \tself.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()\n\t\n\tdef insert_cls(self, x, cls_token):\n\t\t\"\"\" Insert CLS token. \"\"\"\n\t\tcls_tokens = cls_token.expand(x.shape[0], -1, -1)\n\t\tx = torch.cat((cls_tokens, x), dim=1)\n\t\treturn x\n\t\n\tdef remove_cls(self, x):\n\t\t\"\"\" Remove CLS token. \"\"\"\n\t\treturn x[:, 1:, :]\n\t\n\tdef forward(self, x0):\n\t\tB = x0.shape[0]\n\t\t\n\t\t\n\t\t# Serial blocks 1.\n\t\tx1, (H1, W1) = self.patch_embed1(x0)\n\t\tcls = self.cls_token1#torch.zeros_like(self.cls_token1)#self.cls_token1\n\t\tx1 = self.insert_cls(x1, cls)\n\t\tfor blk in self.serial_blocks1:\n\t\t\tx1 = blk(x1, size=(H1, W1))\n\t\tx1_nocls = self.remove_cls(x1)\n\t\tx1_nocls = x1_nocls.reshape(B, H1, W1, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\n\t\t# Serial blocks 2.\n\t\tx2, (H2, W2) = self.patch_embed2(x1_nocls)\n\t\tcls = self.cls_token2# torch.zeros_like(self.cls_token2)#self.cls_token2#\n\t\tx2 = self.insert_cls(x2,cls)\n\t\tfor blk in self.serial_blocks2:\n\t\t\tx2 = blk(x2, size=(H2, W2))\n\t\tx2_nocls = self.remove_cls(x2)\n\t\tx2_nocls = x2_nocls.reshape(B, H2, W2, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\n\t\t# Serial blocks 3.\n\t\tx3, (H3, W3) = self.patch_embed3(x2_nocls)\n\t\tcls = self.cls_token3#torch.zeros_like(self.cls_token3)# self.cls_token3\n\t\tx3 = self.insert_cls(x3, cls)\n\t\tfor blk in self.serial_blocks3:\n\t\t\tx3 = blk(x3, size=(H3, W3))\n\t\tx3_nocls = self.remove_cls(x3)\n\t\tx3_nocls = x3_nocls.reshape(B, H3, W3, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\n\t\t# Serial blocks 4.\n\t\tx4, (H4, W4) = self.patch_embed4(x3_nocls)\n\t\tcls = self.cls_token4#torch.zeros_like(self.cls_token4)#self.cls_token4\n\t\tx4 = self.insert_cls(x4, cls)\n\t\tfor blk in self.serial_blocks4:\n\t\t\tx4 = blk(x4, size=(H4, W4))\n\t\tx4_nocls = self.remove_cls(x4)\n\t\tx4_nocls = x4_nocls.reshape(B, H4, W4, -1).permute(0, 3, 1, 2).contiguous()\n\n\t\t# Serial blocks 5.\n\t\tx5, (H5, W5) = self.patch_embed5(x4_nocls)\n\t\tcls = self.cls_token5#torch.zeros_like(self.cls_token4)#self.cls_token4\n\t\tx5 = self.insert_cls(x5, cls)\n\t\tfor blk in self.serial_blocks5:\n\t\t\tx5 = blk(x5, size=(H5, W5))\n\t\tx5_nocls = self.remove_cls(x5)\n\t\tx5_nocls = x5_nocls.reshape(B, H5, W5, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\n\t\t# Only serial blocks: Early return.\n\t\tif self.parallel_depth == 0:\n\t\t\tx1_nocls = self.out_norm[0](x1_nocls)\n\t\t\tx2_nocls = self.out_norm[1](x2_nocls)\n\t\t\tx3_nocls = self.out_norm[2](x3_nocls)\n\t\t\tx4_nocls = self.out_norm[3](x4_nocls)\n\t\t\treturn [x1_nocls,x2_nocls,x3_nocls,x4_nocls]\n\t\t \n\t \n\t\t\t\n\t\t\n\t\t# Parallel blocks.\n\t\tfor blk in self.parallel_blocks:\n\t\t\tx1, x2, x3, x4,x5 = blk(x1, x2, x3, x4,x5, sizes=[(H1, W1), (H2, W2), (H3, W3), (H4, W4), (H5, W5)])\n\t\t# pdb.set_trace()\n\t\t# remove cls and return feature for seg\n\t\tif self.return_interm_layers:       # Return intermediate features for down-stream tasks (e.g. Deformable DETR and Detectron2).\n\t\t\tfeat_out = {}\n\t\t\tif 'x1_nocls' in self.out_features:\n\t\t\t\tx1_nocls = self.remove_cls(x1)\n\t\t\t\tx1_nocls = x1_nocls.reshape(B, H1, W1, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\t\tfeat_out['x1_nocls'] = x1_nocls\n\t\t\tif 'x2_nocls' in self.out_features:\n\t\t\t\tx2_nocls = self.remove_cls(x2)\n\t\t\t\tx2_nocls = x2_nocls.reshape(B, H2, W2, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\t\tfeat_out['x2_nocls'] = x2_nocls\n\t\t\tif 'x3_nocls' in self.out_features:\n\t\t\t\tx3_nocls = self.remove_cls(x3)\n\t\t\t\tx3_nocls = x3_nocls.reshape(B, H3, W3, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\t\tfeat_out['x3_nocls'] = x3_nocls\n\t\t\tif 'x4_nocls' in self.out_features:\n\t\t\t\tx4_nocls = self.remove_cls(x4)\n\t\t\t\tx4_nocls = x4_nocls.reshape(B, H4, W4, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\t\tfeat_out['x4_nocls'] = x4_nocls\n\t\t\tif 'x5_nocls' in self.out_features:\n\t\t\t\tx5_nocls = self.remove_cls(x5)\n\t\t\t\tx5_nocls = x5_nocls.reshape(B, H5, W5, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\t\tfeat_out['x5_nocls'] = x5_nocls\n\t\t\tfeat_out = list(feat_out.values())\n\t\t\treturn feat_out\n\t\telse:\n\t\t\tx2 = self.norm2(x2)\n\t\t\tx3 = self.norm3(x3)\n\t\t\tx4 = self.norm4(x4)\n\t\t\tx2_cls = x2[:, :1]              # Shape: [B, 1, C].\n\t\t\tx3_cls = x3[:, :1]\n\t\t\tx4_cls = x4[:, :1]\n\t\t\tmerged_cls = torch.cat((x2_cls, x3_cls, x4_cls), dim=1)       # Shape: [B, 3, C].\n\t\t\tmerged_cls = self.aggregate(merged_cls).squeeze(dim=1)        # Shape: [B, C].\n\t\t\treturn merged_cls\n\t\n \n\n\n\n#CoaT.\n#@register_model\n# def coat_tiny(**kwargs):\n# \tmodel = CoaT(patch_size=4, embed_dims=[152, 152, 152, 152], serial_depths=[2, 2, 2, 2], parallel_depth=6, num_heads=8, mlp_ratios=[4, 4, 4, 4], **kwargs)\n# \tmodel.default_cfg = _cfg_coat()\n# \treturn model\n#\n# #@register_model\n# def coat_mini(**kwargs):\n# \tmodel = CoaT(patch_size=4, embed_dims=[152, 216, 216, 216], serial_depths=[2, 2, 2, 2], parallel_depth=6, num_heads=8, mlp_ratios=[4, 4, 4, 4], **kwargs)\n# \tmodel.default_cfg = _cfg_coat()\n# \treturn model\n#\n# #@register_model\n# def coat_small(**kwargs):\n# \tmodel = CoaT(patch_size=4, embed_dims=[152, 320, 320, 320,320], serial_depths=[2, 2, 2, 2,2], parallel_depth=6, num_heads=8, mlp_ratios=[4,4, 4, 4, 4], **kwargs)\n# \tmodel.default_cfg = _cfg_coat()\n# \treturn model\n# #\n# # CoaT-Lite.\n# #@register_model\n# def coat_lite_tiny(**kwargs):\n# \tmodel = CoaT(patch_size=4, embed_dims=[64, 128, 256, 320], serial_depths=[2, 2, 2, 2], parallel_depth=0, num_heads=8, mlp_ratios=[8, 8, 4, 4], **kwargs)\n# \tmodel.default_cfg = _cfg_coat()\n# \treturn model\n#\n# #@register_model\n# def coat_lite_mini(**kwargs):\n# \tmodel = CoaT(patch_size=4, embed_dims=[64, 128, 320, 512], serial_depths=[2, 2, 2, 2], parallel_depth=0, num_heads=8, mlp_ratios=[8, 8, 4, 4], **kwargs)\n# \tmodel.default_cfg = _cfg_coat()\n# \treturn model\n\n#@register_model\nclass coat_lite_small (CoaT):\n\tdef __init__(self, **kwargs):\n\t\tsuper(coat_lite_small, self).__init__(\n\t        patch_size=4, embed_dims=[64, 128, 320, 512], serial_depths=[2, 2, 2, 2],\n\t\t\tparallel_depth=6, num_heads=8, mlp_ratios=[8,8,4,4], **kwargs)\n\t \n\n#@register_model\nclass coat_lite_medium (CoaT):\n\tdef __init__(self, **kwargs):\n\t\tsuper(coat_lite_medium, self).__init__(\n\t\t\tpatch_size=4, embed_dims=[128, 256, 320, 512],\n\t\t\tserial_depths=[3, 6, 10, 8],\n\t\t    parallel_depth=0, num_heads=8, mlp_ratios=[4, 4, 4, 4],\n\t\t    pretrain = 'coat_lite_medium_384x384_f9129688.pth',\n\t\t\t**kwargs)\n\n#@register_model\nclass coat_parallel_small_plus (CoaT):\n\tdef __init__(self, **kwargs):\n\t\tsuper(coat_parallel_small_plus, self).__init__(\n\t\t\tpatch_size=4, embed_dims=[152, 320, 320, 320,320],\n\t\t\tserial_depths=[2, 2, 2, 2,2],\n\t\t\tparallel_depth=6, num_heads=8, mlp_ratios=[4, 4, 4, 4,4], \n\t\t\tpretrain = 'coat_small_7479cf9b.pth',\n\t\t\t**kwargs)\n\n\n            \n#@register_model\nclass coat_parallel_small (CoaT):\n\tdef __init__(self, **kwargs):\n\t\tsuper(coat_parallel_small, self).__init__(\n\t\t\tpatch_size=4, embed_dims=[152, 320, 320, 320],\n\t\t\tserial_depths=[2, 2, 2, 2],\n\t\t\tparallel_depth=6, num_heads=8, mlp_ratios=[4, 4, 4, 4], \n\t\t\tpretrain = 'coat_small_7479cf9b.pth',\n\t\t\t**kwargs)\n\t\t\t\n\t\t\t\n\t\t\t\n\n\nif 0:\n\tnet = coat_parallel_small()\n\t# print(net)\n\tx = torch.rand(2,3,786,786)\n\tout = net(x)\n\t# pdb.set_trace()\n\tprint([f.shape for f in out])\n\t\n\t# checkpoint = '/root/share1/data/pretrain_model/coat_lite_medium_384x384_f9129688.pth'\n\t# checkpoint = torch.load(checkpoint, map_location=lambda storage, loc: storage)\n\t# print(list(checkpoint.keys()))\n\t# #['epoch', 'arch', 'state_dict', 'optimizer', 'version', 'args', 'amp', 'metric']\n\t# state_dict = checkpoint['model']\n\t# print( net.load_state_dict(state_dict,strict=False) )","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:25:21.163879Z","iopub.execute_input":"2022-09-21T07:25:21.164215Z","iopub.status.idle":"2022-09-21T07:25:21.285912Z","shell.execute_reply.started":"2022-09-21T07:25:21.164181Z","shell.execute_reply":"2022-09-21T07:25:21.284770Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# # https://github.com/mlpc-ucsd/CoaT/blob/main/src/models/coat.py\n\n# \"\"\"\n# CoaT architecture.\n\n# Modified from timm/models/vision_transformer.py\n# \"\"\"\n\n# import torch\n# import torch.nn as nn\n# import torch.nn.functional as F\n\n# from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD\n# from timm.models.layers import DropPath, to_2tuple, trunc_normal_\n# from timm.models.registry import register_model\n\n# from einops import rearrange\n# from functools import partial\n# from torch import nn, einsum\n# import pdb\n\n# # __all__ = [\n# # \t\"coat_tiny\",\n# # \t\"coat_mini\",\n# # \t\"coat_small\",\n# # \t\"coat_lite_tiny\",\n# # \t\"coat_lite_mini\",\n# # \t\"coat_lite_small\"\n# # ]\n\n# #---------------------------------------------\n# #https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py#L15\n# class LayerNorm2d(nn.Module):\n# \tdef __init__(self, dim, eps=1e-6):\n# \t\tsuper().__init__()\n# \t\tself.dim = dim\n# \t\tself.weight = nn.Parameter(torch.ones(dim))\n# \t\tself.bias = nn.Parameter(torch.zeros(dim))\n# \t\tself.eps = eps\n\t\n# \tdef forward(self, x):\n# \t\tbatch_size,C,H,W = x.shape\n# \t\t#assert C==self.dim, 'C=%d, self.dim=%d'%(C,self.dim)\n# \t\t#print('C=%d, self.dim=%d'%(C,self.dim))\n\t\t\n# \t\tu = x.mean(1, keepdim=True)\n# \t\ts = (x - u).pow(2).mean(1, keepdim=True)\n# \t\tx = (x - u) / torch.sqrt(s + self.eps)\n# \t\tx = self.weight[:, None, None] * x + self.bias[:, None, None]\n# \t\treturn x\n# #---------------------------------------------\n\n# def _cfg_coat(url='', **kwargs):\n# \treturn {\n# \t\t'url': url,\n# \t\t'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,\n# \t\t'crop_pct': .9, 'interpolation': 'bicubic',\n# \t\t'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,\n# \t\t'first_conv': 'patch_embed.proj', 'classifier': 'head',\n# \t\t**kwargs\n# \t}\n\n\n# class Mlp(nn.Module):\n# \t\"\"\" Feed-forward network (FFN, a.k.a. MLP) class. \"\"\"\n# \tdef __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n# \t\tsuper().__init__()\n# \t\tout_features = out_features or in_features\n# \t\thidden_features = hidden_features or in_features\n# \t\tself.fc1 = nn.Linear(in_features, hidden_features)\n# \t\tself.act = act_layer()\n# \t\tself.fc2 = nn.Linear(hidden_features, out_features)\n# \t\tself.drop = nn.Dropout(drop)\n\t\n# \tdef forward(self, x):\n# \t\tx = self.fc1(x)\n# \t\tx = self.act(x)\n# \t\tx = self.drop(x)\n# \t\tx = self.fc2(x)\n# \t\tx = self.drop(x)\n# \t\treturn x\n\n\n# class ConvRelPosEnc(nn.Module):\n# \t\"\"\" Convolutional relative position encoding. \"\"\"\n# \tdef __init__(self, Ch, h, window):\n# \t\t\"\"\"\n# \t\tInitialization.\n# \t\t\tCh: Channels per head.\n# \t\t\th: Number of heads.\n# \t\t\twindow: Window size(s) in convolutional relative positional encoding. It can have two forms:\n# \t\t\t\t\t1. An integer of window size, which assigns all attention heads with the same window size in ConvRelPosEnc.\n# \t\t\t\t\t2. A dict mapping window size to #attention head splits (e.g. {window size 1: #attention head split 1, window size 2: #attention head split 2})\n# \t\t\t\t\t   It will apply different window size to the attention head splits.\n# \t\t\"\"\"\n# \t\tsuper().__init__()\n\t\t\n# \t\tif isinstance(window, int):\n# \t\t\twindow = {window: h}                                                         # Set the same window size for all attention heads.\n# \t\t\tself.window = window\n# \t\telif isinstance(window, dict):\n# \t\t\tself.window = window\n# \t\telse:\n# \t\t\traise ValueError()\n\t\t\n# \t\tself.conv_list = nn.ModuleList()\n# \t\tself.head_splits = []\n# \t\tfor cur_window, cur_head_split in window.items():\n# \t\t\tdilation = 1                                                                 # Use dilation=1 at default.\n# \t\t\tpadding_size = (cur_window + (cur_window - 1) * (dilation - 1)) // 2         # Determine padding size. Ref: https://discuss.pytorch.org/t/how-to-keep-the-shape-of-input-and-output-same-when-dilation-conv/14338\n# \t\t\tcur_conv = nn.Conv2d(cur_head_split*Ch, cur_head_split*Ch,\n# \t\t\t                     kernel_size=(cur_window, cur_window),\n# \t\t\t                     padding=(padding_size, padding_size),\n# \t\t\t                     dilation=(dilation, dilation),\n# \t\t\t                     groups=cur_head_split*Ch,\n# \t\t\t                     )\n# \t\t\tself.conv_list.append(cur_conv)\n# \t\t\tself.head_splits.append(cur_head_split)\n# \t\tself.channel_splits = [x*Ch for x in self.head_splits]\n\t\n# \tdef forward(self, q, v, size):\n# \t\tB, h, N, Ch = q.shape\n# \t\tH, W = size\n# \t\tassert N == 1 + H * W\n\t\t\n# \t\t# Convolutional relative position encoding.\n# \t\tq_img = q[:,:,1:,:]                                                              # Shape: [B, h, H*W, Ch].\n# \t\tv_img = v[:,:,1:,:]                                                              # Shape: [B, h, H*W, Ch].\n\t\t\n# \t\tv_img = rearrange(v_img, 'B h (H W) Ch -> B (h Ch) H W', H=H, W=W)               # Shape: [B, h, H*W, Ch] -> [B, h*Ch, H, W].\n# \t\tv_img_list = torch.split(v_img, self.channel_splits, dim=1)                      # Split according to channels.\n# \t\tconv_v_img_list = [conv(x) for conv, x in zip(self.conv_list, v_img_list)]\n# \t\tconv_v_img = torch.cat(conv_v_img_list, dim=1)\n# \t\tconv_v_img = rearrange(conv_v_img, 'B (h Ch) H W -> B h (H W) Ch', h=h)          # Shape: [B, h*Ch, H, W] -> [B, h, H*W, Ch].\n\t\t\n# \t\tEV_hat_img = q_img * conv_v_img\n# \t\tzero = torch.zeros((B, h, 1, Ch), dtype=q.dtype, layout=q.layout, device=q.device)\n# \t\tEV_hat = torch.cat((zero, EV_hat_img), dim=2)                                # Shape: [B, h, N, Ch].\n\t\t\n# \t\treturn EV_hat\n\n\n# class FactorAtt_ConvRelPosEnc(nn.Module):\n# \t\"\"\" Factorized attention with convolutional relative position encoding class. \"\"\"\n# \tdef __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., shared_crpe=None):\n# \t\tsuper().__init__()\n# \t\tself.num_heads = num_heads\n# \t\thead_dim = dim // num_heads\n# \t\tself.scale = qk_scale or head_dim ** -0.5\n\t\t\n# \t\tself.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n# \t\tself.attn_drop = nn.Dropout(attn_drop)                                           # Note: attn_drop is actually not used.\n# \t\tself.proj = nn.Linear(dim, dim)\n# \t\tself.proj_drop = nn.Dropout(proj_drop)\n\t\t\n# \t\t# Shared convolutional relative position encoding.\n# \t\tself.crpe = shared_crpe\n\t\n# \tdef forward(self, x, size):\n# \t\tB, N, C = x.shape\n\t\t\n# \t\t# Generate Q, K, V.\n# \t\tqkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)  # Shape: [3, B, h, N, Ch].\n# \t\tq, k, v = qkv[0], qkv[1], qkv[2]                                                 # Shape: [B, h, N, Ch].\n\t\t\n# \t\t# Factorized attention.\n# \t\tk_softmax = k.softmax(dim=2)                                                     # Softmax on dim N.\n# \t\tk_softmax_T_dot_v = einsum('b h n k, b h n v -> b h k v', k_softmax, v)          # Shape: [B, h, Ch, Ch].\n# \t\tfactor_att        = einsum('b h n k, b h k v -> b h n v', q, k_softmax_T_dot_v)  # Shape: [B, h, N, Ch].\n\t\t\n# \t\t# Convolutional relative position encoding.\n# \t\tcrpe = self.crpe(q, v, size=size)                                                # Shape: [B, h, N, Ch].\n\t\t\n# \t\t# Merge and reshape.\n# \t\tx = self.scale * factor_att + crpe\n# \t\tx = x.transpose(1, 2).reshape(B, N, C)                                           # Shape: [B, h, N, Ch] -> [B, N, h, Ch] -> [B, N, C].\n\t\t\n# \t\t# Output projection.\n# \t\tx = self.proj(x)\n# \t\tx = self.proj_drop(x)\n\t\t\n# \t\treturn x                                                                         # Shape: [B, N, C].\n\n\n# class ConvPosEnc(nn.Module):\n# \t\"\"\" Convolutional Position Encoding.\n# \t\tNote: This module is similar to the conditional position encoding in CPVT.\n# \t\"\"\"\n# \tdef __init__(self, dim, k=3):\n# \t\tsuper(ConvPosEnc, self).__init__()\n# \t\tself.proj = nn.Conv2d(dim, dim, k, 1, k//2, groups=dim)\n\t\n# \tdef forward(self, x, size):\n# \t\tB, N, C = x.shape\n# \t\tH, W = size\n# \t\tassert N == 1 + H * W\n\t\t\n# \t\t# Extract CLS token and image tokens.\n# \t\tcls_token, img_tokens = x[:, :1], x[:, 1:]                                       # Shape: [B, 1, C], [B, H*W, C].\n\t\t\n# \t\t# Depthwise convolution.\n# \t\tfeat = img_tokens.transpose(1, 2).view(B, C, H, W)\n# \t\tx = self.proj(feat) + feat\n# \t\tx = x.flatten(2).transpose(1, 2)\n\t\t\n# \t\t# Combine with CLS token.\n# \t\tx = torch.cat((cls_token, x), dim=1)\n\t\t\n# \t\treturn x\n\n\n# class SerialBlock(nn.Module):\n# \t\"\"\" Serial block class.\n# \t\tNote: In this implementation, each serial block only contains a conv-attention and a FFN (MLP) module. \"\"\"\n# \tdef __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,\n# \t             drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,\n# \t             shared_cpe=None, shared_crpe=None):\n# \t\tsuper().__init__()\n\t\t\n# \t\t# Conv-Attention.\n# \t\tself.cpe = shared_cpe\n\t\t\n# \t\tself.norm1 = norm_layer(dim)\n# \t\tself.factoratt_crpe = FactorAtt_ConvRelPosEnc(\n# \t\t\tdim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n# \t\t\tshared_crpe=shared_crpe)\n# \t\tself.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n\t\t\n# \t\t# MLP.\n# \t\tself.norm2 = norm_layer(dim)\n# \t\tmlp_hidden_dim = int(dim * mlp_ratio)\n# \t\tself.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\t\n# \tdef forward(self, x, size):\n# \t\t# Conv-Attention.\n# \t\tx = self.cpe(x, size)                  # Apply convolutional position encoding.\n# \t\tcur = self.norm1(x)\n# \t\tcur = self.factoratt_crpe(cur, size)   # Apply factorized attention and convolutional relative position encoding.\n# \t\tx = x + self.drop_path(cur)\n\t\t\n# \t\t# MLP.\n# \t\tcur = self.norm2(x)\n# \t\tcur = self.mlp(cur)\n# \t\tx = x + self.drop_path(cur)\n\t\t\n# \t\treturn x\n\n\n# class ParallelBlock(nn.Module):\n# \t\"\"\" Parallel block class. \"\"\"\n# \tdef __init__(self, dims, num_heads, mlp_ratios=[], qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,\n# \t             drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,\n# \t             shared_cpes=None, shared_crpes=None):\n# \t\tsuper().__init__()\n\t\t\n# \t\t# Conv-Attention.\n# \t\tself.cpes = shared_cpes\n\t\t\n# \t\tself.norm12 = norm_layer(dims[1])\n# \t\tself.norm13 = norm_layer(dims[2])\n# \t\tself.norm14 = norm_layer(dims[3])\n# \t\t# self.norm15 = norm_layer(dims[4])\n\n\n# \t\tself.factoratt_crpe2 = FactorAtt_ConvRelPosEnc(\n# \t\t\tdims[1], num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n# \t\t\tshared_crpe=shared_crpes[1]\n# \t\t)\n# \t\tself.factoratt_crpe3 = FactorAtt_ConvRelPosEnc(\n# \t\t\tdims[2], num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n# \t\t\tshared_crpe=shared_crpes[2]\n# \t\t)\n# \t\tself.factoratt_crpe4 = FactorAtt_ConvRelPosEnc(\n# \t\t\tdims[3], num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n# \t\t\tshared_crpe=shared_crpes[3]\n# \t\t)\n# \t\t# self.factoratt_crpe5 = FactorAtt_ConvRelPosEnc(\n# \t\t# \tdims[4], num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n# \t\t# \tshared_crpe=shared_crpes[4]\n# \t\t# )\n\n\n# \t\tself.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n\t\t\n# \t\t# MLP.\n# \t\tself.norm22 = norm_layer(dims[1])\n# \t\tself.norm23 = norm_layer(dims[2])\n# \t\tself.norm24 = norm_layer(dims[3])\n# \t\t# self.norm25 = norm_layer(dims[4])\n\n# \t\tassert dims[1] == dims[2] == dims[3]                           # In parallel block, we assume dimensions are the same and share the linear transformation.\n# \t\tassert mlp_ratios[1] == mlp_ratios[2] == mlp_ratios[3]\n# \t\tmlp_hidden_dim = int(dims[1] * mlp_ratios[1])\n# \t\tself.mlp2 = self.mlp3 = self.mlp4 = Mlp(in_features=dims[1], hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\t\n# \tdef upsample(self, x, output_size, size):\n# \t\t\"\"\" Feature map up-sampling. \"\"\"\n# \t\treturn self.interpolate(x, output_size=output_size, size=size)\n\t\n# \tdef downsample(self, x, output_size, size):\n# \t\t\"\"\" Feature map down-sampling. \"\"\"\n# \t\treturn self.interpolate(x, output_size=output_size, size=size)\n\t\n# \tdef interpolate(self, x, output_size, size):\n# \t\t\"\"\" Feature map interpolation. \"\"\"\n# \t\tB, N, C = x.shape\n# \t\tH, W = size\n# \t\tassert N == 1 + H * W\n\t\t\n# \t\tcls_token  = x[:, :1, :]\n# \t\timg_tokens = x[:, 1:, :]\n\t\t\n# \t\timg_tokens = img_tokens.transpose(1, 2).reshape(B, C, H, W)\n# \t\timg_tokens = F.interpolate(img_tokens, size=output_size, mode='bilinear')  # FIXME: May have alignment issue.\n# \t\timg_tokens = img_tokens.reshape(B, C, -1).transpose(1, 2)\n\t\t\n# \t\tout = torch.cat((cls_token, img_tokens), dim=1)\n\t\t\n# \t\treturn out\n\t\n# \tdef forward(self, x1, x2, x3, x4,sizes):\n# \t\t_, (H2, W2), (H3, W3), (H4, W4) = sizes\n\t\t\n# \t\t# Conv-Attention.\n# \t\tx2 = self.cpes[1](x2, size=(H2, W2))  # Note: x1 is ignored.\n# \t\tx3 = self.cpes[2](x3, size=(H3, W3))\n# \t\tx4 = self.cpes[3](x4, size=(H4, W4))\n\t\t\n# \t\tcur2 = self.norm12(x2)\n# \t\tcur3 = self.norm13(x3)\n# \t\tcur4 = self.norm14(x4)\n# \t\tcur2 = self.factoratt_crpe2(cur2, size=(H2,W2))\n# \t\tcur3 = self.factoratt_crpe3(cur3, size=(H3,W3))\n# \t\tcur4 = self.factoratt_crpe4(cur4, size=(H4,W4))\n# \t\tupsample3_2 = self.upsample(cur3, output_size=(H2,W2), size=(H3,W3))\n# \t\tupsample4_3 = self.upsample(cur4, output_size=(H3,W3), size=(H4,W4))\n# \t\tupsample4_2 = self.upsample(cur4, output_size=(H2,W2), size=(H4,W4))\n# \t\tdownsample2_3 = self.downsample(cur2, output_size=(H3,W3), size=(H2,W2))\n# \t\tdownsample3_4 = self.downsample(cur3, output_size=(H4,W4), size=(H3,W3))\n# \t\tdownsample2_4 = self.downsample(cur2, output_size=(H4,W4), size=(H2,W2))\n# \t\tcur2 = cur2  + upsample3_2   + upsample4_2\n# \t\tcur3 = cur3  + upsample4_3   + downsample2_3\n# \t\tcur4 = cur4  + downsample3_4 + downsample2_4\n# \t\tx2 = x2 + self.drop_path(cur2)\n# \t\tx3 = x3 + self.drop_path(cur3)\n# \t\tx4 = x4 + self.drop_path(cur4)\n\t\t\n# \t\t# MLP.\n# \t\tcur2 = self.norm22(x2)\n# \t\tcur3 = self.norm23(x3)\n# \t\tcur4 = self.norm24(x4)\n# \t\tcur2 = self.mlp2(cur2)\n# \t\tcur3 = self.mlp3(cur3)\n# \t\tcur4 = self.mlp4(cur4)\n# \t\tx2 = x2 + self.drop_path(cur2)\n# \t\tx3 = x3 + self.drop_path(cur3)\n# \t\tx4 = x4 + self.drop_path(cur4)\n\t\t\n# \t\treturn x1, x2, x3, x4\n\n\n# class PatchEmbed(nn.Module):\n# \t\"\"\" Image to Patch Embedding \"\"\"\n# \tdef __init__(self, patch_size=16, in_chans=3, embed_dim=768):\n# \t\tsuper().__init__()\n# \t\tpatch_size = to_2tuple(patch_size)\n\t\t\n# \t\tself.patch_size = patch_size\n# \t\tself.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)\n# \t\tself.norm = nn.LayerNorm(embed_dim)\n\t\n# \tdef forward(self, x):\n# \t\t_, _, H, W = x.shape\n# \t\tout_H, out_W = H // self.patch_size[0], W // self.patch_size[1]\n\t\t\n# \t\tx = self.proj(x).flatten(2).transpose(1, 2)\n# \t\tout = self.norm(x)\n\t\t\n# \t\treturn out, (out_H, out_W)\n\n\n# class CoaT(nn.Module):\n# \t\"\"\" CoaT class. \"\"\"\n# \tdef __init__(self, patch_size=16, in_chans=3, embed_dims=[0, 0, 0, 0],\n# \t             serial_depths=[0, 0, 0, 0], parallel_depth=0,\n# \t             num_heads=0, mlp_ratios=[0, 0, 0, 0], qkv_bias=True, qk_scale=None, drop_rate=0., attn_drop_rate=0.,\n# \t             drop_path_rate=0.,\n# \t             norm_layer=partial(nn.LayerNorm, eps=1e-6),\n# \t             return_interm_layers=True,\n# \t             out_features=['x1_nocls','x2_nocls','x3_nocls','x4_nocls','x5_nocls'],\n# \t             crpe_window={3:2, 5:3, 7:3},\n# \t             pretrain=None,\n# \t             out_norm = nn.Identity, #use nn.Identity, nn.BatchNorm2d, LayerNorm2d\n# \t             **kwargs):\n# \t\tsuper().__init__()\n# \t\tself.return_interm_layers = return_interm_layers\n# \t\tself.pretrain     = pretrain\n# \t\tself.embed_dims   = embed_dims\n# \t\tself.out_features = out_features\n# \t\t#self.num_classes  = num_classes\n\t\t\n# \t\t# Patch embeddings.\n# \t\tself.patch_embed1 = PatchEmbed(patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dims[0])\n# \t\tself.patch_embed2 = PatchEmbed(patch_size=2, in_chans=embed_dims[0], embed_dim=embed_dims[1])\n# \t\tself.patch_embed3 = PatchEmbed(patch_size=2, in_chans=embed_dims[1], embed_dim=embed_dims[2])\n# \t\tself.patch_embed4 = PatchEmbed(patch_size=2, in_chans=embed_dims[2], embed_dim=embed_dims[3])\n# \t\t# self.patch_embed5 = PatchEmbed(patch_size=2, in_chans=embed_dims[3], embed_dim=embed_dims[4])\n\t\t\n# \t\t# Class tokens.\n# \t\tself.cls_token1 = nn.Parameter(torch.zeros(1, 1, embed_dims[0]))\n# \t\tself.cls_token2 = nn.Parameter(torch.zeros(1, 1, embed_dims[1]))\n# \t\tself.cls_token3 = nn.Parameter(torch.zeros(1, 1, embed_dims[2]))\n# \t\tself.cls_token4 = nn.Parameter(torch.zeros(1, 1, embed_dims[3]))\n# \t\t# self.cls_token5 = nn.Parameter(torch.zeros(1, 1, embed_dims[4]))\n\t\t\n# \t\t# Convolutional position encodings.\n# \t\tself.cpe1 = ConvPosEnc(dim=embed_dims[0], k=3)\n# \t\tself.cpe2 = ConvPosEnc(dim=embed_dims[1], k=3)\n# \t\tself.cpe3 = ConvPosEnc(dim=embed_dims[2], k=3)\n# \t\tself.cpe4 = ConvPosEnc(dim=embed_dims[3], k=3)\n# \t\t# self.cpe5 = ConvPosEnc(dim=embed_dims[4], k=3)\n\t\t\n# \t\t# Convolutional relative position encodings.\n# \t\tself.crpe1 = ConvRelPosEnc(Ch=embed_dims[0] // num_heads, h=num_heads, window=crpe_window)\n# \t\tself.crpe2 = ConvRelPosEnc(Ch=embed_dims[1] // num_heads, h=num_heads, window=crpe_window)\n# \t\tself.crpe3 = ConvRelPosEnc(Ch=embed_dims[2] // num_heads, h=num_heads, window=crpe_window)\n# \t\tself.crpe4 = ConvRelPosEnc(Ch=embed_dims[3] // num_heads, h=num_heads, window=crpe_window)\n# \t\t# self.crpe5 = ConvRelPosEnc(Ch=embed_dims[4] // num_heads, h=num_heads, window=crpe_window)\n\n# \t\t# Enable stochastic depth.\n# \t\tdpr = drop_path_rate\n\t\t\n# \t\t# Serial blocks 1.\n# \t\tself.serial_blocks1 = nn.ModuleList([\n# \t\t\tSerialBlock(\n# \t\t\t\tdim=embed_dims[0], num_heads=num_heads, mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale,\n# \t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n# \t\t\t\tshared_cpe=self.cpe1, shared_crpe=self.crpe1\n# \t\t\t)\n# \t\t\tfor _ in range(serial_depths[0])]\n# \t\t)\n\t\t\n# \t\t# Serial blocks 2.\n# \t\tself.serial_blocks2 = nn.ModuleList([\n# \t\t\tSerialBlock(\n# \t\t\t\tdim=embed_dims[1], num_heads=num_heads, mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale,\n# \t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n# \t\t\t\tshared_cpe=self.cpe2, shared_crpe=self.crpe2\n# \t\t\t)\n# \t\t\tfor _ in range(serial_depths[1])]\n# \t\t)\n\t\t\n# \t\t# Serial blocks 3.\n# \t\tself.serial_blocks3 = nn.ModuleList([\n# \t\t\tSerialBlock(\n# \t\t\t\tdim=embed_dims[2], num_heads=num_heads, mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale,\n# \t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n# \t\t\t\tshared_cpe=self.cpe3, shared_crpe=self.crpe3\n# \t\t\t)\n# \t\t\tfor _ in range(serial_depths[2])]\n# \t\t)\n\t\t\n# \t\t# Serial blocks 4.\n# \t\tself.serial_blocks4 = nn.ModuleList([\n# \t\t\tSerialBlock(\n# \t\t\t\tdim=embed_dims[3], num_heads=num_heads, mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale,\n# \t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n# \t\t\t\tshared_cpe=self.cpe4, shared_crpe=self.crpe4\n# \t\t\t)\n# \t\t\tfor _ in range(serial_depths[3])]\n# \t\t)\n\n# \t\t# # Serial blocks 5.\n# \t\t# self.serial_blocks5 = nn.ModuleList([\n# \t\t# \tSerialBlock(\n# \t\t# \t\tdim=embed_dims[4], num_heads=num_heads, mlp_ratio=mlp_ratios[4], qkv_bias=qkv_bias, qk_scale=qk_scale,\n# \t\t# \t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n# \t\t# \t\tshared_cpe=self.cpe4, shared_crpe=self.crpe4\n# \t\t# \t)\n# \t\t# \tfor _ in range(serial_depths[4])]\n# \t\t# )\n\t\t\n# \t\t# Parallel blocks.\n# \t\tself.parallel_depth = parallel_depth\n# \t\tif self.parallel_depth > 0:\n# \t\t\tself.parallel_blocks = nn.ModuleList([\n# \t\t\t\tParallelBlock(\n# \t\t\t\t\tdims=embed_dims, num_heads=num_heads, mlp_ratios=mlp_ratios, qkv_bias=qkv_bias, qk_scale=qk_scale,\n# \t\t\t\t\tdrop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr, norm_layer=norm_layer,\n# \t\t\t\t\tshared_cpes=[self.cpe1, self.cpe2, self.cpe3, self.cpe4],\n# \t\t\t\t\tshared_crpes=[self.crpe1, self.crpe2, self.crpe3, self.crpe4]\n# \t\t\t\t)\n# \t\t\t\tfor _ in range(parallel_depth)]\n# \t\t\t)\n\t\t\n# \t\t# Classification head(s).\n# \t\t# if not self.return_interm_layers:\n# \t\t# \tself.norm1 = norm_layer(embed_dims[0])\n# \t\t# \tself.norm2 = norm_layer(embed_dims[1])\n# \t\t# \tself.norm3 = norm_layer(embed_dims[2])\n# \t\t# \tself.norm4 = norm_layer(embed_dims[3])\n# \t\t#\n# \t\t# \tif self.parallel_depth > 0:                                  # CoaT series: Aggregate features of last three scales for classification.\n# \t\t# \t\tassert embed_dims[1] == embed_dims[2] == embed_dims[3]\n# \t\t# \t\tself.aggregate = torch.nn.Conv1d(in_channels=3, out_channels=1, kernel_size=1)\n# \t\t# \t\tself.head = nn.Linear(embed_dims[3], num_classes)\n# \t\t# \telse:\n# \t\t# \t\tself.head = nn.Linear(embed_dims[3], num_classes)        # CoaT-Lite series: Use feature of last scale for classification.\n# \t\t#---\n# \t\t# add a norm layer for each output\n# \t\tself.out_norm = nn.ModuleList(\n# \t\t\t[ out_norm(embed_dims[i]) for i in range(4)]\n# \t\t)\n\t\t\n# \t\t# Initialize weights.\n# \t\ttrunc_normal_(self.cls_token1, std=.02)\n# \t\ttrunc_normal_(self.cls_token2, std=.02)\n# \t\ttrunc_normal_(self.cls_token3, std=.02)\n# \t\ttrunc_normal_(self.cls_token4, std=.02)\n# \t\t# trunc_normal_(self.cls_token5, std=.02)\n# \t\tself.apply(self._init_weights)\n\t\n# \tdef _init_weights(self, m):\n# \t\tif isinstance(m, nn.Linear):\n# \t\t\ttrunc_normal_(m.weight, std=.02)\n# \t\t\tif isinstance(m, nn.Linear) and m.bias is not None:\n# \t\t\t\tnn.init.constant_(m.bias, 0)\n# \t\telif isinstance(m, nn.LayerNorm):\n# \t\t\tnn.init.constant_(m.bias, 0)\n# \t\t\tnn.init.constant_(m.weight, 1.0)\n\t\n# \t@torch.jit.ignore\n# \tdef no_weight_decay(self):\n# \t\treturn {'cls_token1', 'cls_token2', 'cls_token3', 'cls_token4'}\n\t\n# \t# def get_classifier(self):\n# \t# \treturn self.head\n\t\n# \t# def reset_classifier(self, num_classes, global_pool=''):\n# \t# \tself.num_classes = num_classes\n# \t# \tself.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()\n\t\n# \tdef insert_cls(self, x, cls_token):\n# \t\t\"\"\" Insert CLS token. \"\"\"\n# \t\tcls_tokens = cls_token.expand(x.shape[0], -1, -1)\n# \t\tx = torch.cat((cls_tokens, x), dim=1)\n# \t\treturn x\n\t\n# \tdef remove_cls(self, x):\n# \t\t\"\"\" Remove CLS token. \"\"\"\n# \t\treturn x[:, 1:, :]\n\t\n# \tdef forward(self, x0):\n# \t\tB = x0.shape[0]\n\t\t\n\t\t\n# \t\t# Serial blocks 1.\n# \t\tx1, (H1, W1) = self.patch_embed1(x0)\n# \t\tcls = self.cls_token1#torch.zeros_like(self.cls_token1)#self.cls_token1\n# \t\tx1 = self.insert_cls(x1, cls)\n# \t\tfor blk in self.serial_blocks1:\n# \t\t\tx1 = blk(x1, size=(H1, W1))\n# \t\tx1_nocls = self.remove_cls(x1)\n# \t\tx1_nocls = x1_nocls.reshape(B, H1, W1, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\n# \t\t# Serial blocks 2.\n# \t\tx2, (H2, W2) = self.patch_embed2(x1_nocls)\n# \t\tcls = self.cls_token2# torch.zeros_like(self.cls_token2)#self.cls_token2#\n# \t\tx2 = self.insert_cls(x2,cls)\n# \t\tfor blk in self.serial_blocks2:\n# \t\t\tx2 = blk(x2, size=(H2, W2))\n# \t\tx2_nocls = self.remove_cls(x2)\n# \t\tx2_nocls = x2_nocls.reshape(B, H2, W2, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\n# \t\t# Serial blocks 3.\n# \t\tx3, (H3, W3) = self.patch_embed3(x2_nocls)\n# \t\tcls = self.cls_token3#torch.zeros_like(self.cls_token3)# self.cls_token3\n# \t\tx3 = self.insert_cls(x3, cls)\n# \t\tfor blk in self.serial_blocks3:\n# \t\t\tx3 = blk(x3, size=(H3, W3))\n# \t\tx3_nocls = self.remove_cls(x3)\n# \t\tx3_nocls = x3_nocls.reshape(B, H3, W3, -1).permute(0, 3, 1, 2).contiguous()\n\t\t\n# \t\t# Serial blocks 4.\n# \t\tx4, (H4, W4) = self.patch_embed4(x3_nocls)\n# \t\tcls = self.cls_token4#torch.zeros_like(self.cls_token4)#self.cls_token4\n# \t\tx4 = self.insert_cls(x4, cls)\n# \t\tfor blk in self.serial_blocks4:\n# \t\t\tx4 = blk(x4, size=(H4, W4))\n# \t\tx4_nocls = self.remove_cls(x4)\n# \t\tx4_nocls = x4_nocls.reshape(B, H4, W4, -1).permute(0, 3, 1, 2).contiguous()\n\n\n\t\t\n# \t\t# Only serial blocks: Early return.\n# \t\t# if self.parallel_depth == 0:\n# \t\t# \tx1_nocls = self.out_norm[0](x1_nocls)\n# \t\t# \tx2_nocls = self.out_norm[1](x2_nocls)\n# \t\t# \tx3_nocls = self.out_norm[2](x3_nocls)\n# \t\t# \tx4_nocls = self.out_norm[3](x4_nocls)\n# \t\t# \treturn [x1_nocls,x2_nocls,x3_nocls,x4_nocls]\n\t\t \n\t \n\t\t\t\n\t\t\n# \t\t# Parallel blocks.\n# \t\tfor blk in self.parallel_blocks:\n# \t\t\tx1, x2, x3, x4 = blk(x1, x2, x3, x4, sizes=[(H1, W1), (H2, W2), (H3, W3), (H4, W4)])\n# \t\t# pdb.set_trace()\n# \t\t# remove cls and return feature for seg\n# \t\tif self.return_interm_layers:       # Return intermediate features for down-stream tasks (e.g. Deformable DETR and Detectron2).\n# \t\t\tfeat_out = {}\n# \t\t\tif 'x1_nocls' in self.out_features:\n# \t\t\t\tx1_nocls = self.remove_cls(x1)\n# \t\t\t\tx1_nocls = x1_nocls.reshape(B, H1, W1, -1).permute(0, 3, 1, 2).contiguous()\n# \t\t\t\tfeat_out['x1_nocls'] = x1_nocls\n# \t\t\tif 'x2_nocls' in self.out_features:\n# \t\t\t\tx2_nocls = self.remove_cls(x2)\n# \t\t\t\tx2_nocls = x2_nocls.reshape(B, H2, W2, -1).permute(0, 3, 1, 2).contiguous()\n# \t\t\t\tfeat_out['x2_nocls'] = x2_nocls\n# \t\t\tif 'x3_nocls' in self.out_features:\n# \t\t\t\tx3_nocls = self.remove_cls(x3)\n# \t\t\t\tx3_nocls = x3_nocls.reshape(B, H3, W3, -1).permute(0, 3, 1, 2).contiguous()\n# \t\t\t\tfeat_out['x3_nocls'] = x3_nocls\n# \t\t\tif 'x4_nocls' in self.out_features:\n# \t\t\t\tx4_nocls = self.remove_cls(x4)\n# \t\t\t\tx4_nocls = x4_nocls.reshape(B, H4, W4, -1).permute(0, 3, 1, 2).contiguous()\n# \t\t\t\tfeat_out['x4_nocls'] = x4_nocls\n\n# \t\t\tfeat_out = list(feat_out.values())\n# \t\t\treturn feat_out\n# \t\telse:\n# \t\t\tx2 = self.norm2(x2)\n# \t\t\tx3 = self.norm3(x3)\n# \t\t\tx4 = self.norm4(x4)\n# \t\t\tx2_cls = x2[:, :1]              # Shape: [B, 1, C].\n# \t\t\tx3_cls = x3[:, :1]\n# \t\t\tx4_cls = x4[:, :1]\n# \t\t\tmerged_cls = torch.cat((x2_cls, x3_cls, x4_cls), dim=1)       # Shape: [B, 3, C].\n# \t\t\tmerged_cls = self.aggregate(merged_cls).squeeze(dim=1)        # Shape: [B, C].\n# \t\t\treturn merged_cls\n\t\n \n\n\n\n# #CoaT.\n# #@register_model\n# # def coat_tiny(**kwargs):\n# # \tmodel = CoaT(patch_size=4, embed_dims=[152, 152, 152, 152], serial_depths=[2, 2, 2, 2], parallel_depth=6, num_heads=8, mlp_ratios=[4, 4, 4, 4], **kwargs)\n# # \tmodel.default_cfg = _cfg_coat()\n# # \treturn model\n# #\n# # #@register_model\n# # def coat_mini(**kwargs):\n# # \tmodel = CoaT(patch_size=4, embed_dims=[152, 216, 216, 216], serial_depths=[2, 2, 2, 2], parallel_depth=6, num_heads=8, mlp_ratios=[4, 4, 4, 4], **kwargs)\n# # \tmodel.default_cfg = _cfg_coat()\n# # \treturn model\n# #\n# # #@register_model\n# # def coat_small(**kwargs):\n# # \tmodel = CoaT(patch_size=4, embed_dims=[152, 320, 320, 320,320], serial_depths=[2, 2, 2, 2,2], parallel_depth=6, num_heads=8, mlp_ratios=[4,4, 4, 4, 4], **kwargs)\n# # \tmodel.default_cfg = _cfg_coat()\n# # \treturn model\n# # #\n# # # CoaT-Lite.\n# # #@register_model\n# # def coat_lite_tiny(**kwargs):\n# # \tmodel = CoaT(patch_size=4, embed_dims=[64, 128, 256, 320], serial_depths=[2, 2, 2, 2], parallel_depth=0, num_heads=8, mlp_ratios=[8, 8, 4, 4], **kwargs)\n# # \tmodel.default_cfg = _cfg_coat()\n# # \treturn model\n# #\n# # #@register_model\n# # def coat_lite_mini(**kwargs):\n# # \tmodel = CoaT(patch_size=4, embed_dims=[64, 128, 320, 512], serial_depths=[2, 2, 2, 2], parallel_depth=0, num_heads=8, mlp_ratios=[8, 8, 4, 4], **kwargs)\n# # \tmodel.default_cfg = _cfg_coat()\n# # \treturn model\n\n# #@register_model\n# class coat_lite_small (CoaT):\n# \tdef __init__(self, **kwargs):\n# \t\tsuper(coat_lite_small, self).__init__(\n# \t        patch_size=4, embed_dims=[64, 128, 320, 512], serial_depths=[2, 2, 2, 2],\n# \t\t\tparallel_depth=0, num_heads=8, mlp_ratios=[8,8,4,4], **kwargs)\n\t \n\n# #@register_model\n# class coat_lite_medium (CoaT):\n# \tdef __init__(self, **kwargs):\n# \t\tsuper(coat_lite_medium, self).__init__(\n# \t\t\tpatch_size=4, embed_dims=[128, 256, 320, 512],\n# \t\t\tserial_depths=[3, 6, 10, 8],\n# \t\t    parallel_depth=0, num_heads=8, mlp_ratios=[4, 4, 4, 4],\n# \t\t    pretrain = 'coat_lite_medium_384x384_f9129688.pth',\n# \t\t\t**kwargs)\n\n# #@register_model\n# class coat_parallel_small (CoaT):\n# \tdef __init__(self, **kwargs):\n# \t\tsuper(coat_parallel_small, self).__init__(\n# \t\t\tpatch_size=4, embed_dims=[152, 320, 320, 320],    #embed_dims=[128, 320, 320, 320],\n#  \t\t\tserial_depths=[2, 2, 2, 2],\n# \t\t\tparallel_depth=6, num_heads=8, mlp_ratios=[4, 4, 4, 4], \n# \t\t\tpretrain = 'coat_small_7479cf9b.pth',\n# \t\t\t**kwargs)\n\t\t\t\n\t\t\t\n\t\t\t\n\n\n# if 0:\n# \tnet = coat_parallel_small()\n# \t# print(net)\n# \tx = torch.rand(2,3,786,786)\n# \tout = net(x)\n# \t# pdb.set_trace()\n# \tprint([f.shape for f in out])\n\t\n# \t# checkpoint = '/root/share1/data/pretrain_model/coat_lite_medium_384x384_f9129688.pth'\n# \t# checkpoint = torch.load(checkpoint, map_location=lambda storage, loc: storage)\n# \t# print(list(checkpoint.keys()))\n# \t# #['epoch', 'arch', 'state_dict', 'optimizer', 'version', 'args', 'amp', 'metric']\n# \t# state_dict = checkpoint['model']\n# \t# print( net.load_state_dict(state_dict,strict=False) )\n\t","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:25:21.287674Z","iopub.execute_input":"2022-09-21T07:25:21.288096Z","iopub.status.idle":"2022-09-21T07:25:21.316845Z","shell.execute_reply.started":"2022-09-21T07:25:21.288059Z","shell.execute_reply":"2022-09-21T07:25:21.315924Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nimport pdb\n\n\n#################################################################\n\nclass RGB(nn.Module):\n\tIMAGE_RGB_MEAN = [0.485, 0.456, 0.406]  # [0.5, 0.5, 0.5]\n\tIMAGE_RGB_STD = [0.229, 0.224, 0.225]  # [0.5, 0.5, 0.5]\n\t\n\tdef __init__(self, ):\n\t\tsuper(RGB, self).__init__()\n\t\tself.register_buffer('mean', torch.zeros(1, 3, 1, 1))\n\t\tself.register_buffer('std', torch.ones(1, 3, 1, 1))\n\t\tself.mean.data = torch.FloatTensor(self.IMAGE_RGB_MEAN).view(self.mean.shape)\n\t\tself.std.data = torch.FloatTensor(self.IMAGE_RGB_STD).view(self.std.shape)\n\t\n\tdef forward(self, x):\n\t\tx = (x - self.mean) / self.std\n\t\treturn x\n\n\n\ndef init_weight(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        #nn.init.xavier_uniform_(m.weight, gain=1)\n        #nn.init.xavier_normal_(m.weight, gain=1)\n        #nn.init.kaiming_uniform_(m.weight, mode='fan_in', nonlinearity='relu')\n        nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')\n        #nn.init.orthogonal_(m.weight, gain=1)\n        if m.bias is not None:\n            m.bias.data.zero_()\n    elif classname.find('Batch') != -1:\n        m.weight.data.normal_(1,0.02)\n        m.bias.data.zero_()\n    elif classname.find('Linear') != -1:\n        nn.init.orthogonal_(m.weight, gain=1)\n        if m.bias is not None:\n            m.bias.data.zero_()\n    elif classname.find('Embedding') != -1:\n        nn.init.orthogonal_(m.weight, gain=1)\n\n\n\nclass Net(nn.Module):\n\t# def load_pretrain( self,):\n\t# \tpath =  '/home/jupyter/share/WRQ/Hubmap/input/coat_lite_medium/coat_lite_medium_384x384_f9129688.pth'\n\t# \tcheckpoint = torch.load(path, map_location=lambda storage, loc: storage)\n\t# \tself.state_dict = checkpoint['model']\n\t# \tself.encoder.load_state_dict(checkpoint,strict=False)\n\n\n\tdef __init__(self,\n\t             \n\t             decoder=daformer_conv1x1,   \n\t\t\t\t \n\t             encoder_cfg={},\n\t             decoder_cfg={},\n\t             ):  # decoder = daformer_conv3x3,   for coat-medium\n\t\tsuper(Net, self).__init__()\n\t\tdecoder_dim = decoder_cfg.get('decoder_dim', 320)\n\t\t\n\t\t# ----\n\t\tself.output_type = ['inference', 'loss']\n\t\t\n\n\t\tself.rgb = RGB()\n    \n\t\t# encoder=coat_lite_medium()\n\t\tself.encoder=coat_parallel_small_plus()\n\n\n\t\t\n\t\t\n\t\tencoder_dim = self.encoder.embed_dims\n\t\t# pdb.set_trace()\n\t\t# [64, 128, 320, 512]\t\n\t\t# [128, 256, 320, 512],\n\t\t\n\t\tself.decoder = decoder(\n\t\t\tencoder_dim=encoder_dim,\n\t\t\tdecoder_dim=decoder_dim,\n\t\t)\n\t\tself.logit = nn.Sequential(\n\t\t\tnn.Conv2d(decoder_dim, 1, kernel_size=1),\n\t\t)\n\n\t\tself.aux = nn.ModuleList([\n            nn.Conv2d(decoder_dim, 1, kernel_size=1, padding=0) for i in range(4)\n        ])\n\t\tself.avgpool = nn.AdaptiveAvgPool2d((1,1))\n\t\t# self.cls_head = nn.Linear(320,5,bias = False)\n\t\tself.cls_head = nn.Sequential(\n\t\t\t\t\t\t\tnn.BatchNorm1d(320).apply(init_weight),\n\t\t\t\t\t\t\tnn.Linear(320, 128 ).apply(init_weight),\n\t\t\t\t\t\t\t\n\t\t\t\t\t\t\tnn.ReLU(inplace=True),\n\t\t\t\t\t\t\tnn.BatchNorm1d(128).apply(init_weight),\n\t\t\t\t\t\t\tnn.Linear(128, 5).apply(init_weight)\n\t\t\t\t\t\t)\n\t \n\t\n\n\tdef forward(self, batch):\n\t\t\n\t\tx = batch\n\t\t\n\t\tnum_class = 5\n\n\n\t\tx = self.rgb(x)\n\t\t\n\t\tB, C, H, W = x.shape\n\t\t# pdb.set_trace()\n        \n\t\tencoder = self.encoder(x)\n\t\t# print([f.shape for f in encoder])\n\t\t# cls_feature = encoder[-1]\n\t\t# cls_feature = self.avgpool(cls_feature)\n\t\t# cls_feature = cls_feature.view(cls_feature.size(0), -1)\n\t\t# cls_feature = self.cls_head(cls_feature)        # (batch,num_class)\n\n\n\n   \n\n\n\n\t\t# pdb.set_trace()\n\t\tlast, decoder = self.decoder(encoder)\n\t\t# pdb.set_trace()\n\t\tlogit = self.logit(last)\n\t\tlogit = F.interpolate(logit, size=None, scale_factor=4, mode='bilinear', align_corners=False)\n\t\t# print(logit.shape)\n\t\t# pdb.set_trace()\n\t\toutput = {}\n\t\tif 'loss' in self.output_type:\n\t\t\toutput['bce_loss'] = F.binary_cross_entropy_with_logits(logit,batch['mask'])\n\t\t\t# pdb.set_trace()\n\t\t\t# output[\"label_loss\"] = F.nll_loss(F.log_softmax(cls_feature,dim=1), label)\n\t\t\t# output[\"label_loss\"] = nn.CrossEntropyLoss()(cls_feature,label)\n\t\t\tfor i in range(4):\n\t\t\t\toutput['aux%d_loss'%i] = criterion_aux_loss(self.aux[i](decoder[i]),batch['mask'])\n\t\tif 'inference' in self.output_type:\n\t\t\tprobability_from_logit = torch.sigmoid(logit)\n\t\t\toutput['probability'] = probability_from_logit\n\t\t\n\t\treturn output\n\ndef criterion_aux_loss(logit, mask):\n    mask = F.interpolate(mask,size=logit.shape[-2:], mode='nearest')\n    loss = F.binary_cross_entropy_with_logits(logit,mask)\n    return loss\n\n \n\n# def one_hot(label, depth = 5):\n\n#     y = torch.zeros(label.size(0), depth)\n#     idx = torch.tensor(label).view(-1, 1).cuda()\n#     y.scatter_(dim = 1, index = idx, value = 1)\n#     return y\n \n# def criterion_multi_binary_cross_entropy(logit, mask, organ):\n# \tlogit = F.interpolate(logit, size=None, scale_factor=4, mode='bilinear', align_corners=False)\n# \tbatch_size, C, H, W = logit.shape\n\t\n# \tlabel = mask.long() * organ.reshape(batch_size, 1, 1, 1)\n# \tonehot = torch.zeros((batch_size, num_organ + 1, H, W)).to(mask)\n# \tonehot = onehot.scatter(1, label, 1)\n# \t# onehot[:,0] = 1-onehot[:,0]\n\t\n# \tloss = F.binary_cross_entropy_with_logits(logit, onehot)\n# \treturn loss\n\n\ndef run_check_net():\n\tbatch_size = 2\n\timage_size = 768\n\t\n\t# ---\n\tbatch = {\n\t\t'image': torch.from_numpy(np.random.uniform(-1, 1, (batch_size, 3, image_size, image_size))).float(),\n\t\t'mask': torch.from_numpy(np.random.choice(2, (batch_size, 1, image_size, image_size))).float(),\n\t\t'organ': torch.from_numpy(np.random.choice(5, (batch_size, 1))).long(),\n\t}\n\tbatch = {k: v for k, v in batch.items()}\n\t\n\tnet = Net()\n\t# net.load_pretrain()\n\twith torch.no_grad():\n\t\twith torch.cuda.amp.autocast(enabled=True):\n\t\t\toutput = net(batch)\n\t\n\tprint('batch')\n\tfor k, v in batch.items():\n\t\tprint('%32s :' % k, v.shape)\n# \t\n\tprint('output')\n\tfor k, v in output.items():\n\t\tif 'loss' not in k:\n\t\t\tprint('%32s :' % k, v.shape)\n\tfor k, v in output.items():\n\t\tif 'loss' in k:\n\t\t\tprint('%32s :' % k, v.item())\n\n\n\n# # # main #################################################################\n# if __name__ == '__main__':\n# \trun_check_net()","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:26:06.092668Z","iopub.execute_input":"2022-09-21T07:26:06.093182Z","iopub.status.idle":"2022-09-21T07:26:06.133157Z","shell.execute_reply.started":"2022-09-21T07:26:06.093138Z","shell.execute_reply":"2022-09-21T07:26:06.132087Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"### import glob\nimport copy\nclass CFG:\n#     step1: hyper-parameter\n    seed = 42 \n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n    num_worker = 0 # if debug\n    data_path = \"../input/hubmap-organ-segmentation\"\n#     ckpt_paths =  \"../input/coatsmall4level768swakaimingfold3\"   \n    ckpt_paths =  \"../input/coat-5layer\"\n#     ckpt_paths =  \"../input/coat-5layer/3-45-score-0.8601-loss-0.0865.pth\"\n    # step2: data\n    n_fold = 5\n    img_size = [1024, 1024]   #  img_size = [768, 768]\n    train_bs = 1\n    valid_bs = train_bs * 2\n    # step3: model\n    backbone = 'swin+upernet'\n    num_classes = 1\n    # step4: optimizer\n    epoch = 12\n    lr = 1e-3\n    wd = 1e-5\n    lr_drop = 8\n    # step5: infer\n#     thr = 0.3\n    thr = {\n        \"Hubmap\":{\n            \"kidney\" : 0.35,\n            \"prostate\":0.26,\n            \"largeintestine\":0.32,\n            \"spleen\":0.28,\n            \"lung\":0.08,  \n        },\n        \n         \"HPA\":{\n            \"kidney\" : 0.4,\n            \"prostate\":0.4,\n            \"largeintestine\":0.4,\n            \"spleen\":0.4,\n            \"lung\":0.1,\n             \n         },}\n    tta = True\n\n###############################################################\n##### Data prepare is from original dataset shape which is beetween (0,1)\n###############################################################\n\ndef read_tiff(path, scale=None, verbose=0): \n    image = tifffile.imread(path)\n    if len(image.shape) == 5:\n        image = image.squeeze().transpose(1, 2, 0)\n    \n    if verbose:\n        print(f\"[{path}] Image shape: {image.shape}\")\n    \n    if scale:\n        new_size = (image.shape[1] // scale, image.shape[0] // cale)\n        image = cv2.resize(image, new_size)\n        if verbose:\n            print(f\"[{path}] Resized Image shape: {image.shape}\")\n        \n    mx = np.max(image)\n    image = image.astype(np.float32)\n    if mx:\n        image /= mx # scale image to [0, 1]\n    return image\n        \n#https://www.kaggle.com/bguberfain/memory-aware-rle-encoding\n#with transposed mask\n\ndef rle_encode_less_memory(img):\n    #the image should be transposed\n    pixels = img.T.flatten()\n    \n    # This simplified method requires first and last pixel to be zero\n    pixels[0] = 0\n    pixels[-1] = 0\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    runs[1::2] -= runs[::2]\n    \n    return ' '.join(str(x) for x in runs)\n\ndef rle_encode(img):\n    '''\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = img.T.flatten()\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return ' '.join(str(x) for x in runs)\n\n# Ref: https://www.kaggle.com/code/paulorzp/rle-functions-run-lenght-encode-decode/script\ndef rle_decode(mask_rle, shape):\n    '''\n    mask_rle: run-length as string formated (start length)\n    shape: (height,width) of array to return \n    Returns numpy array, 1 - mask, 0 - background\n    '''\n    s = mask_rle.split()\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)\n    for lo, hi in zip(starts, ends):\n        img[lo:hi] = 1\n    return img.reshape(shape).T  # Needed to align to RLE direction\n\n\ndef build_transforms(CFG):\n    data_transforms = {\n        \"valid_test\": A.Compose([\n            A.Resize(*CFG.img_size, interpolation=cv2.INTER_NEAREST),\n            ], p=1.0)\n        }\n    return data_transforms\n\n\nclass build_dataset(Dataset):\n    def __init__(self, df, label=True, transforms=None):\n        self.df = df\n        self.label = label\n        self.transforms = transforms\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, index):\n        \n        img_path = self.df.loc[index, 'image_path']\n        img_height = self.df.loc[index, 'img_height']\n        img_width = self.df.loc[index, 'img_width']\n        organs = self.df.loc[index, 'organ']\n        id_ = self.df.loc[index, 'id']\n        img = read_tiff(img_path)\n        sours = self.df.loc[index,'data_source']\n        \n        if self.label:\n            rle_mask = self.df.loc[index, 'rle']\n            mask = rle_decode(rle_mask, (img_height, img_width))\n            # pdb.set_trace() \n            if self.transforms:\n                data = self.transforms(image=img, mask=mask)\n                img  = data['image']\n                mask  = data['mask']\n            \n            mask = np.expand_dims(mask, axis=0)\n            img = np.transpose(img, (2, 0, 1))\n            # mask = np.transpose(mask, (2, 0, 1))\n            \n            return torch.tensor(img), torch.tensor(mask)\n        \n        else:    # resize for infer\n            if self.transforms:\n                data = self.transforms(image=img)\n                img  = data['image']\n                \n            img = np.transpose(img, (2, 0, 1))   #(c, h, w)\n            return torch.tensor(img), img_height, img_width,id_,organs,sours\n#             return torch.tensor(img), img_height, img_width, id_,organ\n\n\n# trained model for infer\n\ndef build_model(ckpt_path,test_flag = True):\n\n    model = Net()\n    model.output_type = [\"inference\"]\n    print(ckpt_path)\n#     model.load_state_dict(copy.deepcopy(torch.load(ckpt_path)[\"state_dict\"])(),strict=False\n    model.load_state_dict(torch.load(ckpt_path),strict=False)\n    model.float()\n    model.eval()\n    model.to(CFG.device)\n    return model\n\n\n# single fold\n# def build_model(CFG, test_flag = True):\n\n#     model = Net()\n#     model.output_type = [\"inference\"]\n# #     pdb.set_trace()\n# #     model.load_state_dict(copy.deepcopy(torch.load(CFG.ckpt_paths)[\"state_dict\"])(),strict=False)\n# #     pdb.set_trace()\n#     model.load_state_dict(torch.load(CFG.ckpt_paths),strict=False)\n# #     pdb.set_trace()\n#     model.float()\n#     model.eval()\n#     model.to(CFG.device)\n#     return model\n\n# def build_model(CFG, test_flag = True):\n\n#     model = Net()\n#     model.output_type = [\"inference\"]\n# #     pdb.set_trace()\n# #     model.load_state_dict(torch.load(CFG.ckpt_paths)[\"state_dict\"],strict=False)\n#     model.load_state_dict(copy.deepcopy(torch.load(CFG.ckpt_paths)[\"state_dict\"])(),strict=False)\n#     model.float()\n#     model.eval()\n#     model.to(CFG.device)\n#     return model\n\n\n\n\n# @torch.no_grad()\n# def test_one_epoch(test_loader,CFG):\n#     pred_ids = []\n#     pred_rles = []\n    \n#     pbar = tqdm(enumerate(test_loader), total=len(test_loader), desc='Test: ')\n#     for _, (images, heights, widths, ids,organs,sours) in pbar:\n# #         pdb.set_trace()\n#         images  = images.to(CFG.device, dtype=torch.float) # [b, c, w, h]\n#         size = images.size()\n#         masks = torch.zeros((size[0], CFG.num_classes, size[2], size[3]), device=CFG.device, dtype=torch.float32) # [b, c, w, h]\n        \n#         ############################################\n#         ##### >>>>>>> cross validation infer <<<<<<\n#         ############################################\n#         # for sub_ckpt_path in ckpt_paths:\n#         model = build_model(CFG, test_flag=True)\n#         # model.load_state_dict(torch.load(sub_ckpt_path))\n#         # model.load_state_dict(torch.load(CFG.ckpt_paths))\n#         # model.eval()\n#         output = model(images) # [b, c, w, h]\n#         y_pred = output[\"probability\"]\n        \n     \n#         masks = y_pred\n# #         organ = organs[0]\n# #         sour = sours[0]\n        \n# #         thr = CFG.thr[sour][organ]\n#         if organs[0] == \"lung\":\n#             CFG.thr = 0.1\n# #         pdb.set_trace()\n#         masks = (masks.permute((0, 2, 3, 1))>CFG.thr).to(torch.uint8).cpu().detach().numpy() # [n, h, w, c]\n# #         masks = (masks.permute((0, 2, 3, 1))>thr).to(torch.uint8).cpu().detach().numpy() # [n, h, w, c]\n \n#         for idx in range(masks.shape[0]):\n#             height = heights[idx].item()\n#             width = widths[idx].item()\n#             id_ = ids[idx].item()\n#             msk = cv2.resize(masks[idx].squeeze(), dsize=(width, height), interpolation=cv2.INTER_NEAREST)\n#             rle = rle_encode_less_memory(msk)\n#             pred_rles.append(rle)\n#             pred_ids.append(id_)\n    \n#     return pred_ids, pred_rles,msk\n\n\n@torch.no_grad()\ndef test_one_epoch(test_loader,ckpt_paths,CFG):\n    pred_ids = []\n    pred_rles = []\n    \n    pbar = tqdm(enumerate(test_loader), total=len(test_loader), desc='Test: ')\n    for _, (images, heights, widths, ids,organs,sours) in pbar:\n\n        images  = images.to(CFG.device, dtype=torch.float) # [b, c, w, h]\n        size = images.size()\n        masks = torch.zeros((size[0], CFG.num_classes, size[2], size[3]), device=CFG.device, dtype=torch.float32) # [b, c, w, h]\n        \n        ###########################################\n        #### >>>>>>> cross validation infer <<<<<<\n        ###########################################\n#         pdb.set_trace()\n        for sub_ckpt_path in ckpt_paths:\n#             if '212' in sub_ckpt_path:continue\n#             if '45' in sub_ckpt_path:continue\n            if '288' in sub_ckpt_path:continue\n            if '331' in sub_ckpt_path:continue\n\n\n            model = build_model(sub_ckpt_path, test_flag=True)\n          \n            output = model(images) # [b, c, w, h]\n            y_pred = output[\"probability\"]\n\n\n            \n            masks += y_pred\n\n            if CFG.tta:\n\n#                 flips = [[-1],[-2],[-2,-1]]\n                flips = [[-1],[-2]]\n                for f in flips:\n                   \n                    images_f = torch.flip(images, f)\n                    y_pred = model(images_f)[\"probability\"] # [b, c, w, h]\n                    y_pred = torch.flip(y_pred, f)\n                    masks += y_pred\n        \n        if CFG.tta:\n            total_ckpt_paths = len(ckpt_paths) * 3\n#             pdb.set_trace()\n            \n        else:\n            total_ckpt_paths = len(ckpt_paths)\n        \n        masks /= total_ckpt_paths   \n#         pdb.set_trace()\n       \n        \n        organ = organs[0]\n        sour = sours[0]\n        \n        thr = CFG.thr[sour][organ]\n#         if organs[0] == \"lung\":\n#             CFG.thr = 0.1\n#         masks = (masks.permute((0, 2, 3, 1))>CFG.thr).to(torch.uint8).cpu().detach().numpy() # [n, h, w, c]\n        masks = (masks.permute((0, 2, 3, 1))>thr).to(torch.uint8).cpu().detach().numpy() # [n, h, w, c]\n             \n \n\n        for idx in range(masks.shape[0]):\n            height = heights[idx].item()\n            width = widths[idx].item()\n            id_ = ids[idx].item()\n            msk = cv2.resize(masks[idx].squeeze(), dsize=(width, height), interpolation=cv2.INTER_NEAREST)\n            rle = rle_encode_less_memory(msk)\n            pred_rles.append(rle)\n            pred_ids.append(id_)\n    \n    return pred_ids, pred_rles,msk \n\n\n\nif __name__ == '__main__':\n\n    ###############################################################\n    ##### part0: read data  which contains test.csv file and Corresponding image\n    ###############################################################\n    df = pd.read_csv(os.path.join(CFG.data_path, \"test.csv\"))\n    df['image_path'] = df['id'].apply(lambda x: os.path.join(CFG.data_path, 'test_images', str(x) + '.tiff'))\n\n    data_transforms = build_transforms(CFG)\n    test_dataset = build_dataset(df, label=False, transforms=data_transforms['valid_test'])\n    \n    test_loader  = DataLoader(test_dataset, batch_size=1, num_workers=0, shuffle=False, pin_memory=False)\n    ckpt_paths  = glob.glob(CFG.ckpt_paths + \"/*.pth\")\n#     pdb.set_trace()\n    print(\"ckptpath:{}\",CFG.ckpt_paths)\n#     pdb.set_trace()\n#     ckpt_paths  = [\"../input/coatdaformer/00005520.model.pth\",\n#                    \"../input/coatdaformer/00005796.model.pth\"\n#                     ]\n    ###############################################################\n    ##### >>>>>>> step2: infer <<<<<<\n    ###############################################################\n    # attention: change the corresponding upload path to kaggle.\n    # ckpt_paths  = glob(f'{ckpt_path}/best*')\n    # assert len(ckpt_paths) == n_fold, \"ckpt path error!\"\n#     pred_ids, pred_rles,msk = test_one_epoch(test_loader,CFG)\n    \n    pred_ids, pred_rles,msk = test_one_epoch(test_loader, ckpt_paths,CFG)\n# #     pdb.set_trace()\n    msk.astype(np.float32)\n    plt.imshow(msk*255)\n#     pdb.set_trace()\n#     # plt.show()\n   \n    # plt.savefig(\"mask\")\n    \n\n\n\n\n    ###############################################################\n    #####  IF Local Debug:\n    ###############################################################\n   \n        \n    \n    ###############################################################\n    ##### step3: submit\n    ###############################################################\n    pred_df = pd.DataFrame({\n        \"id\":pred_ids,\n        \"rle\":pred_rles\n        \n    })\n    pred_df.to_csv('submission.csv',index=False)","metadata":{"execution":{"iopub.status.busy":"2022-09-21T07:39:14.014153Z","iopub.execute_input":"2022-09-21T07:39:14.014534Z","iopub.status.idle":"2022-09-21T07:39:18.878324Z","shell.execute_reply.started":"2022-09-21T07:39:14.014502Z","shell.execute_reply":"2022-09-21T07:39:18.877332Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}