vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/aspanformer/src/ASpanFormer/aspan_module/transformer.py

@@ -0,0 +1,244 @@
+import copy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .attention import FullAttention, HierachicalAttention ,layernorm2d
+
+
+class messageLayer_ini(nn.Module):
+
+    def __init__(self, d_model, d_flow,d_value, nhead):
+        super().__init__()
+        super(messageLayer_ini, self).__init__()
+
+        self.d_model = d_model
+        self.d_flow = d_flow
+        self.d_value=d_value
+        self.nhead = nhead
+        self.attention = FullAttention(d_model,nhead)
+
+        self.q_proj = nn.Conv1d(d_model, d_model, kernel_size=1,bias=False)
+        self.k_proj = nn.Conv1d(d_model, d_model, kernel_size=1,bias=False)
+        self.v_proj = nn.Conv1d(d_value, d_model, kernel_size=1,bias=False)
+        self.merge_head=nn.Conv1d(d_model,d_model,kernel_size=1,bias=False)
+
+        self.merge_f= self.merge_f = nn.Sequential(
+            nn.Conv2d(d_model*2, d_model*2, kernel_size=1, bias=False),
+            nn.ReLU(True),
+            nn.Conv2d(d_model*2, d_model, kernel_size=1, bias=False),
+        )
+
+        self.norm1 = layernorm2d(d_model)
+        self.norm2 = layernorm2d(d_model)
+
+
+    def forward(self, x0, x1,pos0,pos1,mask0=None,mask1=None):
+        #x1,x2: b*d*L
+        x0,x1=self.update(x0,x1,pos1,mask0,mask1),\
+                self.update(x1,x0,pos0,mask1,mask0)
+        return x0,x1
+
+
+    def update(self,f0,f1,pos1,mask0,mask1):
+        """
+        Args:
+            f0: [N, D, H, W]
+            f1: [N, D, H, W]
+        Returns:
+            f0_new: (N, d, h, w)
+        """
+        bs,h,w=f0.shape[0],f0.shape[2],f0.shape[3]
+
+        f0_flatten,f1_flatten=f0.view(bs,self.d_model,-1),f1.view(bs,self.d_model,-1)
+        pos1_flatten=pos1.view(bs,self.d_value-self.d_model,-1)
+        f1_flatten_v=torch.cat([f1_flatten,pos1_flatten],dim=1)
+
+        queries,keys=self.q_proj(f0_flatten),self.k_proj(f1_flatten)
+        values=self.v_proj(f1_flatten_v).view(bs,self.nhead,self.d_model//self.nhead,-1)
+        
+        queried_values=self.attention(queries,keys,values,mask0,mask1)
+        msg=self.merge_head(queried_values).view(bs,-1,h,w)
+        msg=self.norm2(self.merge_f(torch.cat([f0,self.norm1(msg)],dim=1)))
+        return f0+msg
+
+
+
+class messageLayer_gla(nn.Module):
+
+    def __init__(self,d_model,d_flow,d_value,
+                    nhead,radius_scale,nsample,update_flow=True):
+        super().__init__()
+        self.d_model = d_model
+        self.d_flow=d_flow
+        self.d_value=d_value
+        self.nhead = nhead
+        self.radius_scale=radius_scale
+        self.update_flow=update_flow
+        self.flow_decoder=nn.Sequential(
+                    nn.Conv1d(d_flow, d_flow//2, kernel_size=1, bias=False),
+                    nn.ReLU(True),
+                    nn.Conv1d(d_flow//2, 4, kernel_size=1, bias=False))
+        self.attention=HierachicalAttention(d_model,nhead,nsample,radius_scale)
+
+        self.q_proj = nn.Conv1d(d_model, d_model, kernel_size=1,bias=False)
+        self.k_proj = nn.Conv1d(d_model, d_model, kernel_size=1,bias=False)
+        self.v_proj = nn.Conv1d(d_value, d_model, kernel_size=1,bias=False)
+
+        d_extra=d_flow if update_flow else 0
+        self.merge_f=nn.Sequential(
+                     nn.Conv2d(d_model*2+d_extra, d_model+d_flow, kernel_size=1, bias=False),
+                     nn.ReLU(True),
+                     nn.Conv2d(d_model+d_flow, d_model+d_extra, kernel_size=3,padding=1, bias=False),
+                )
+        self.norm1 = layernorm2d(d_model)
+        self.norm2 = layernorm2d(d_model+d_extra)
+
+    def forward(self, x0, x1, flow_feature0,flow_feature1,pos0,pos1,mask0=None,mask1=None,ds0=[4,4],ds1=[4,4]):
+        """
+        Args:
+            x0 (torch.Tensor): [B, C, H, W]
+            x1 (torch.Tensor): [B, C, H, W]
+            flow_feature0 (torch.Tensor): [B, C', H, W]
+            flow_feature1 (torch.Tensor): [B, C', H, W]
+        """
+        flow0,flow1=self.decode_flow(flow_feature0,flow_feature1.shape[2:]),self.decode_flow(flow_feature1,flow_feature0.shape[2:])
+        x0_new,flow_feature0_new=self.update(x0,x1,flow0.detach(),flow_feature0,pos1,mask0,mask1,ds0,ds1)
+        x1_new,flow_feature1_new=self.update(x1,x0,flow1.detach(),flow_feature1,pos0,mask1,mask0,ds1,ds0)
+        return x0_new,x1_new,flow_feature0_new,flow_feature1_new,flow0,flow1
+
+    def update(self,x0,x1,flow0,flow_feature0,pos1,mask0,mask1,ds0,ds1):
+        bs=x0.shape[0]
+        queries,keys=self.q_proj(x0.view(bs,self.d_model,-1)),self.k_proj(x1.view(bs,self.d_model,-1))
+        x1_pos=torch.cat([x1,pos1],dim=1)
+        values=self.v_proj(x1_pos.view(bs,self.d_value,-1))
+        msg=self.attention(queries,keys,values,flow0,x0.shape[2:],x1.shape[2:],mask0,mask1,ds0,ds1)
+
+        if self.update_flow:
+            update_feature=torch.cat([x0,flow_feature0],dim=1)
+        else:
+            update_feature=x0
+        msg=self.norm2(self.merge_f(torch.cat([update_feature,self.norm1(msg)],dim=1)))
+        update_feature=update_feature+msg
+
+        x0_new,flow_feature0_new=update_feature[:,:self.d_model],update_feature[:,self.d_model:]
+        return x0_new,flow_feature0_new
+
+    def decode_flow(self,flow_feature,kshape):
+        bs,h,w=flow_feature.shape[0],flow_feature.shape[2],flow_feature.shape[3]
+        scale_factor=torch.tensor([kshape[1],kshape[0]]).to(flow_feature.device)[None,None,None]
+        flow=self.flow_decoder(flow_feature.view(bs,-1,h*w)).permute(0,2,1).view(bs,h,w,4)
+        flow_coordinates=torch.sigmoid(flow[:,:,:,:2])*scale_factor
+        flow_var=flow[:,:,:,2:]
+        flow=torch.cat([flow_coordinates,flow_var],dim=-1) #B*H*W*4
+        return flow
+
+
+class flow_initializer(nn.Module):
+
+    def __init__(self, dim, dim_flow, nhead, layer_num):
+        super().__init__()
+        self.layer_num= layer_num
+        self.dim = dim
+        self.dim_flow = dim_flow
+
+        encoder_layer = messageLayer_ini(
+            dim ,dim_flow,dim+dim_flow , nhead)
+        self.layers_coarse = nn.ModuleList(
+            [copy.deepcopy(encoder_layer) for _ in range(layer_num)])
+        self.decoupler = nn.Conv2d(
+                self.dim, self.dim+self.dim_flow, kernel_size=1)
+        self.up_merge = nn.Conv2d(2*dim, dim, kernel_size=1)
+
+    def forward(self, feat0, feat1,pos0,pos1,mask0=None,mask1=None,ds0=[4,4],ds1=[4,4]):
+        # feat0: [B, C, H0, W0]
+        # feat1: [B, C, H1, W1]
+        # use low-res MHA to initialize flow feature
+        bs = feat0.size(0)
+        h0,w0,h1,w1=feat0.shape[2],feat0.shape[3],feat1.shape[2],feat1.shape[3]
+
+        # coarse level
+        sub_feat0, sub_feat1 = F.avg_pool2d(feat0, ds0, stride=ds0), \
+                            F.avg_pool2d(feat1, ds1, stride=ds1)
+
+        sub_pos0,sub_pos1=F.avg_pool2d(pos0, ds0, stride=ds0), \
+                            F.avg_pool2d(pos1, ds1, stride=ds1)
+    
+        if mask0 is not None:
+            mask0,mask1=-F.max_pool2d(-mask0.view(bs,1,h0,w0),ds0,stride=ds0).view(bs,-1),\
+                        -F.max_pool2d(-mask1.view(bs,1,h1,w1),ds1,stride=ds1).view(bs,-1)
+        
+        for layer in self.layers_coarse:
+            sub_feat0, sub_feat1 = layer(sub_feat0, sub_feat1,sub_pos0,sub_pos1,mask0,mask1)
+        # decouple flow and visual features
+        decoupled_feature0, decoupled_feature1 = self.decoupler(sub_feat0),self.decoupler(sub_feat1) 
+
+        sub_feat0, sub_flow_feature0 = decoupled_feature0[:,:self.dim], decoupled_feature0[:, self.dim:]
+        sub_feat1, sub_flow_feature1 = decoupled_feature1[:,:self.dim], decoupled_feature1[:, self.dim:]
+        update_feat0, flow_feature0 = F.upsample(sub_feat0, scale_factor=ds0, mode='bilinear'),\
+                                        F.upsample(sub_flow_feature0, scale_factor=ds0, mode='bilinear')
+        update_feat1, flow_feature1 = F.upsample(sub_feat1, scale_factor=ds1, mode='bilinear'),\
+                                        F.upsample(sub_flow_feature1, scale_factor=ds1, mode='bilinear')
+        
+        feat0 = feat0+self.up_merge(torch.cat([feat0, update_feat0], dim=1))
+        feat1 = feat1+self.up_merge(torch.cat([feat1, update_feat1], dim=1))
+    
+        return feat0,feat1,flow_feature0,flow_feature1 #b*c*h*w
+
+
+class LocalFeatureTransformer_Flow(nn.Module):
+    """A Local Feature Transformer (LoFTR) module."""
+
+    def __init__(self, config):
+        super(LocalFeatureTransformer_Flow, self).__init__()
+
+        self.config = config
+        self.d_model = config['d_model']
+        self.nhead = config['nhead']
+
+        self.pos_transform=nn.Conv2d(config['d_model'],config['d_flow'],kernel_size=1,bias=False)
+        self.ini_layer = flow_initializer(self.d_model, config['d_flow'], config['nhead'],config['ini_layer_num'])
+        
+        encoder_layer = messageLayer_gla(
+            config['d_model'], config['d_flow'], config['d_flow']+config['d_model'], config['nhead'],config['radius_scale'],config['nsample'])
+        encoder_layer_last=messageLayer_gla(
+            config['d_model'], config['d_flow'], config['d_flow']+config['d_model'], config['nhead'],config['radius_scale'],config['nsample'],update_flow=False)
+        self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(config['layer_num']-1)]+[encoder_layer_last])
+        self._reset_parameters()
+        
+    def _reset_parameters(self):
+        for name,p in self.named_parameters():
+            if 'temp' in name or 'sample_offset' in name:
+                continue
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, feat0, feat1,pos0,pos1,mask0=None,mask1=None,ds0=[4,4],ds1=[4,4]):
+        """
+        Args:
+            feat0 (torch.Tensor): [N, C, H, W]
+            feat1 (torch.Tensor): [N, C, H, W]
+            pos1,pos2:  [N, C, H, W]
+        Outputs:
+            feat0: [N,-1,C]
+            feat1: [N,-1,C]
+            flow_list: [L,N,H,W,4]*1(2)
+        """
+        bs = feat0.size(0)
+        
+        pos0,pos1=self.pos_transform(pos0),self.pos_transform(pos1)
+        pos0,pos1=pos0.expand(bs,-1,-1,-1),pos1.expand(bs,-1,-1,-1)
+        assert self.d_model == feat0.size(
+            1), "the feature number of src and transformer must be equal"
+       
+        flow_list=[[],[]]# [px,py,sx,sy] 
+        if mask0 is not None:
+            mask0,mask1=mask0[:,None].float(),mask1[:,None].float()
+        feat0,feat1, flow_feature0, flow_feature1 = self.ini_layer(feat0, feat1,pos0,pos1,mask0,mask1,ds0,ds1)
+        for layer in self.layers:
+            feat0,feat1,flow_feature0,flow_feature1,flow0,flow1=layer(feat0,feat1,flow_feature0,flow_feature1,pos0,pos1,mask0,mask1,ds0,ds1)
+            flow_list[0].append(flow0)
+            flow_list[1].append(flow1)
+        flow_list[0]=torch.stack(flow_list[0],dim=0)
+        flow_list[1]=torch.stack(flow_list[1],dim=0)
+        feat0, feat1 = feat0.permute(0, 2, 3, 1).view(bs, -1, self.d_model), feat1.permute(0, 2, 3, 1).view(bs, -1, self.d_model)
+        return feat0, feat1, flow_list

vismatch/third_party/aspanformer/src/ASpanFormer/aspanformer.py

@@ -0,0 +1,152 @@
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from einops.einops import rearrange
+
+from .backbone import build_backbone
+from .utils.position_encoding import PositionEncodingSine
+from .aspan_module import LocalFeatureTransformer_Flow, LocalFeatureTransformer, FinePreprocess
+from .utils.coarse_matching import CoarseMatching
+from .utils.fine_matching import FineMatching
+
+from typing import Dict
+from kornia.core import Module, Tensor
+
+
+class ASpanFormer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        # Misc
+        self.config = config
+
+        # Modules
+        self.backbone = build_backbone(config)
+        self.pos_encoding = PositionEncodingSine(
+            config['coarse']['d_model'],pre_scaling=[config['coarse']['train_res'],config['coarse']['test_res']])
+        self.loftr_coarse = LocalFeatureTransformer_Flow(config['coarse'])
+        self.coarse_matching = CoarseMatching(config['match_coarse'])
+        self.fine_preprocess = FinePreprocess(config)
+        self.loftr_fine = LocalFeatureTransformer(config["fine"])
+        self.fine_matching = FineMatching()
+        self.coarsest_level=config['coarse']['coarsest_level']
+
+    def forward(self, data, online_resize=False):
+        """ 
+        Update:
+            data (dict): {
+                'image0': (torch.Tensor): (N, 1, H, W)
+                'image1': (torch.Tensor): (N, 1, H, W)
+                'mask0'(optional) : (torch.Tensor): (N, H, W) '0' indicates a padded position
+                'mask1'(optional) : (torch.Tensor): (N, H, W)
+            }
+        """
+        if online_resize:
+            assert data['image0'].shape[0]==1 and data['image1'].shape[1]==1
+            self.resize_input(data,self.config['coarse']['train_res'])
+        else:
+            data['pos_scale0'],data['pos_scale1']=None,None
+
+        # 1. Local Feature CNN
+        data.update({
+            'bs': data['image0'].size(0),
+            'hw0_i': data['image0'].shape[2:], 'hw1_i': data['image1'].shape[2:]
+        })
+        
+        if data['hw0_i'] == data['hw1_i']:  # faster & better BN convergence
+            feats_c, feats_f = self.backbone(
+                torch.cat([data['image0'], data['image1']], dim=0))
+            (feat_c0, feat_c1), (feat_f0, feat_f1) = feats_c.split(
+                data['bs']), feats_f.split(data['bs'])
+        else:  # handle different input shapes
+            (feat_c0, feat_f0), (feat_c1, feat_f1) = self.backbone(
+                data['image0']), self.backbone(data['image1'])
+
+        data.update({
+            'hw0_c': feat_c0.shape[2:], 'hw1_c': feat_c1.shape[2:],
+            'hw0_f': feat_f0.shape[2:], 'hw1_f': feat_f1.shape[2:]
+        })
+
+        # 2. coarse-level loftr module
+        # add featmap with positional encoding, then flatten it to sequence [N, HW, C]
+        [feat_c0, pos_encoding0], [feat_c1, pos_encoding1] = self.pos_encoding(feat_c0,data['pos_scale0']), self.pos_encoding(feat_c1,data['pos_scale1'])
+        feat_c0 = rearrange(feat_c0, 'n c h w -> n c h w ')
+        feat_c1 = rearrange(feat_c1, 'n c h w -> n c h w ')
+
+        #TODO:adjust ds 
+        ds0=[int(data['hw0_c'][0]/self.coarsest_level[0]),int(data['hw0_c'][1]/self.coarsest_level[1])]
+        ds1=[int(data['hw1_c'][0]/self.coarsest_level[0]),int(data['hw1_c'][1]/self.coarsest_level[1])]
+        if online_resize:
+            ds0,ds1=[4,4],[4,4]
+
+        mask_c0 = mask_c1 = None  # mask is useful in training
+        if 'mask0' in data:
+            mask_c0, mask_c1 = data['mask0'].flatten(
+                -2), data['mask1'].flatten(-2)
+        feat_c0, feat_c1, flow_list = self.loftr_coarse(
+            feat_c0, feat_c1,pos_encoding0,pos_encoding1,mask_c0,mask_c1,ds0,ds1)
+
+        # 3. match coarse-level and register predicted offset
+        self.coarse_matching(feat_c0, feat_c1, flow_list,data,
+                             mask_c0=mask_c0, mask_c1=mask_c1)
+
+        # 4. fine-level refinement
+        feat_f0_unfold, feat_f1_unfold = self.fine_preprocess(
+            feat_f0, feat_f1, feat_c0, feat_c1, data)
+        if feat_f0_unfold.size(0) != 0:  # at least one coarse level predicted
+            feat_f0_unfold, feat_f1_unfold = self.loftr_fine(
+                feat_f0_unfold, feat_f1_unfold)
+
+        # 5. match fine-level
+        self.fine_matching(feat_f0_unfold, feat_f1_unfold, data)
+
+        # 6. resize match coordinates back to input resolution
+        if online_resize:
+            data['mkpts0_f']*=data['online_resize_scale0']
+            data['mkpts1_f']*=data['online_resize_scale1']
+
+        rename_keys: Dict[str, str] = {
+            "mkpts0_f": 'keypoints0',
+            "mkpts1_f": 'keypoints1',
+            "mconf": 'confidence',
+            "b_ids": 'batch_indexes',
+        }
+        out: Dict[str, Tensor] = {}
+        for k, v in rename_keys.items():
+            _d = data[k]
+            if isinstance(_d, Tensor):
+                out[v] = _d
+            else:
+                raise TypeError(f'Expected Tensor for item `{k}`. Gotcha {type(_d)}')
+        return out
+        
+    def load_state_dict(self, state_dict, *args, **kwargs):
+        for k in list(state_dict.keys()):
+            if k.startswith('matcher.'):
+                if 'sample_offset' in k:
+                    state_dict.pop(k)
+                else:
+                    state_dict[k.replace('matcher.', '', 1)] = state_dict.pop(k)
+        return super().load_state_dict(state_dict, *args, **kwargs)
+    
+    def resize_input(self,data,train_res,df=32):
+        device = data['image0'].device
+        h0,w0,h1,w1=data['image0'].shape[2],data['image0'].shape[3],data['image1'].shape[2],data['image1'].shape[3]
+        data['image0'],data['image1']=self.resize_df(data['image0'],df),self.resize_df(data['image1'],df)
+        
+        if len(train_res)==1:
+            train_res_h=train_res_w=train_res
+        else:
+            train_res_h,train_res_w=train_res[0],train_res[1]
+        data['pos_scale0'],data['pos_scale1']=[train_res_h/data['image0'].shape[2],train_res_w/data['image0'].shape[3]],\
+                                  [train_res_h/data['image1'].shape[2],train_res_w/data['image1'].shape[3]] 
+        data['online_resize_scale0'],data['online_resize_scale1']=torch.tensor([w0/data['image0'].shape[3],h0/data['image0'].shape[2]])[None].to(device),\
+                                                                    torch.tensor([w1/data['image1'].shape[3],h1/data['image1'].shape[2]])[None].to(device)
+
+    def resize_df(self,image,df=32):
+        h,w=image.shape[2],image.shape[3]
+        h_new,w_new=h//df*df,w//df*df
+        if h!=h_new or w!=w_new:
+            img_new=transforms.Resize([h_new,w_new]).forward(image)
+        else:
+            img_new=image
+        return img_new

vismatch/third_party/aspanformer/src/ASpanFormer/backbone/__init__.py

@@ -0,0 +1,11 @@
+from .resnet_fpn import ResNetFPN_8_2, ResNetFPN_16_4
+
+
+def build_backbone(config):
+    if config['backbone_type'] == 'ResNetFPN':
+        if config['resolution'] == (8, 2):
+            return ResNetFPN_8_2(config['resnetfpn'])
+        elif config['resolution'] == (16, 4):
+            return ResNetFPN_16_4(config['resnetfpn'])
+    else:
+        raise ValueError(f"LOFTR.BACKBONE_TYPE {config['backbone_type']} not supported.")

vismatch/third_party/aspanformer/src/ASpanFormer/backbone/resnet_fpn.py

@@ -0,0 +1,199 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution without padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False)
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_planes, planes, stride=1):
+        super().__init__()
+        self.conv1 = conv3x3(in_planes, planes, stride)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+
+        if stride == 1:
+            self.downsample = None
+        else:
+            self.downsample = nn.Sequential(
+                conv1x1(in_planes, planes, stride=stride),
+                nn.BatchNorm2d(planes)
+            )
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.bn1(self.conv1(y)))
+        y = self.bn2(self.conv2(y))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+
+class ResNetFPN_8_2(nn.Module):
+    """
+    ResNet+FPN, output resolution are 1/8 and 1/2.
+    Each block has 2 layers.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        # Config
+        block = BasicBlock
+        initial_dim = config['initial_dim']
+        block_dims = config['block_dims']
+
+        # Class Variable
+        self.in_planes = initial_dim
+
+        # Networks
+        self.conv1 = nn.Conv2d(1, initial_dim, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(initial_dim)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.layer1 = self._make_layer(block, block_dims[0], stride=1)  # 1/2
+        self.layer2 = self._make_layer(block, block_dims[1], stride=2)  # 1/4
+        self.layer3 = self._make_layer(block, block_dims[2], stride=2)  # 1/8
+
+        # 3. FPN upsample
+        self.layer3_outconv = conv1x1(block_dims[2], block_dims[2])
+        self.layer2_outconv = conv1x1(block_dims[1], block_dims[2])
+        self.layer2_outconv2 = nn.Sequential(
+            conv3x3(block_dims[2], block_dims[2]),
+            nn.BatchNorm2d(block_dims[2]),
+            nn.LeakyReLU(),
+            conv3x3(block_dims[2], block_dims[1]),
+        )
+        self.layer1_outconv = conv1x1(block_dims[0], block_dims[1])
+        self.layer1_outconv2 = nn.Sequential(
+            conv3x3(block_dims[1], block_dims[1]),
+            nn.BatchNorm2d(block_dims[1]),
+            nn.LeakyReLU(),
+            conv3x3(block_dims[1], block_dims[0]),
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, dim, stride=1):
+        layer1 = block(self.in_planes, dim, stride=stride)
+        layer2 = block(dim, dim, stride=1)
+        layers = (layer1, layer2)
+
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        # ResNet Backbone
+        x0 = self.relu(self.bn1(self.conv1(x)))
+        x1 = self.layer1(x0)  # 1/2
+        x2 = self.layer2(x1)  # 1/4
+        x3 = self.layer3(x2)  # 1/8
+
+        # FPN
+        x3_out = self.layer3_outconv(x3)
+
+        x3_out_2x = F.interpolate(x3_out, scale_factor=2., mode='bilinear', align_corners=True)
+        x2_out = self.layer2_outconv(x2)
+        x2_out = self.layer2_outconv2(x2_out+x3_out_2x)
+
+        x2_out_2x = F.interpolate(x2_out, scale_factor=2., mode='bilinear', align_corners=True)
+        x1_out = self.layer1_outconv(x1)
+        x1_out = self.layer1_outconv2(x1_out+x2_out_2x)
+
+        return [x3_out, x1_out]
+
+
+class ResNetFPN_16_4(nn.Module):
+    """
+    ResNet+FPN, output resolution are 1/16 and 1/4.
+    Each block has 2 layers.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        # Config
+        block = BasicBlock
+        initial_dim = config['initial_dim']
+        block_dims = config['block_dims']
+
+        # Class Variable
+        self.in_planes = initial_dim
+
+        # Networks
+        self.conv1 = nn.Conv2d(1, initial_dim, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(initial_dim)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.layer1 = self._make_layer(block, block_dims[0], stride=1)  # 1/2
+        self.layer2 = self._make_layer(block, block_dims[1], stride=2)  # 1/4
+        self.layer3 = self._make_layer(block, block_dims[2], stride=2)  # 1/8
+        self.layer4 = self._make_layer(block, block_dims[3], stride=2)  # 1/16
+
+        # 3. FPN upsample
+        self.layer4_outconv = conv1x1(block_dims[3], block_dims[3])
+        self.layer3_outconv = conv1x1(block_dims[2], block_dims[3])
+        self.layer3_outconv2 = nn.Sequential(
+            conv3x3(block_dims[3], block_dims[3]),
+            nn.BatchNorm2d(block_dims[3]),
+            nn.LeakyReLU(),
+            conv3x3(block_dims[3], block_dims[2]),
+        )
+
+        self.layer2_outconv = conv1x1(block_dims[1], block_dims[2])
+        self.layer2_outconv2 = nn.Sequential(
+            conv3x3(block_dims[2], block_dims[2]),
+            nn.BatchNorm2d(block_dims[2]),
+            nn.LeakyReLU(),
+            conv3x3(block_dims[2], block_dims[1]),
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, dim, stride=1):
+        layer1 = block(self.in_planes, dim, stride=stride)
+        layer2 = block(dim, dim, stride=1)
+        layers = (layer1, layer2)
+
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        # ResNet Backbone
+        x0 = self.relu(self.bn1(self.conv1(x)))
+        x1 = self.layer1(x0)  # 1/2
+        x2 = self.layer2(x1)  # 1/4
+        x3 = self.layer3(x2)  # 1/8
+        x4 = self.layer4(x3)  # 1/16
+
+        # FPN
+        x4_out = self.layer4_outconv(x4)
+
+        x4_out_2x = F.interpolate(x4_out, scale_factor=2., mode='bilinear', align_corners=True)
+        x3_out = self.layer3_outconv(x3)
+        x3_out = self.layer3_outconv2(x3_out+x4_out_2x)
+
+        x3_out_2x = F.interpolate(x3_out, scale_factor=2., mode='bilinear', align_corners=True)
+        x2_out = self.layer2_outconv(x2)
+        x2_out = self.layer2_outconv2(x2_out+x3_out_2x)
+
+        return [x4_out, x2_out]