File size: 3,686 Bytes
b13b124 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
import math
from .helpers import load_pretrained
from .layers import DropPath, to_2tuple, trunc_normal_
from ..builder import HEADS
from .decode_head import BaseDecodeHead
from ..backbones.vit import Block
from mmcv.cnn import build_norm_layer
class MLAHead(nn.Module):
def __init__(self, mla_channels=256, mlahead_channels=128, norm_cfg=None):
super(MLAHead, self).__init__()
self.head2 = nn.Sequential(nn.Conv2d(mla_channels, mlahead_channels, 3, padding=1, bias=False),
build_norm_layer(norm_cfg, mlahead_channels)[1], nn.ReLU(),
nn.Conv2d(mlahead_channels, mlahead_channels, 3, padding=1, bias=False),
build_norm_layer(norm_cfg, mlahead_channels)[1], nn.ReLU())
self.head3 = nn.Sequential(nn.Conv2d(mla_channels, mlahead_channels, 3, padding=1, bias=False),
build_norm_layer(norm_cfg, mlahead_channels)[1], nn.ReLU(),
nn.Conv2d(mlahead_channels, mlahead_channels, 3, padding=1, bias=False),
build_norm_layer(norm_cfg, mlahead_channels)[1], nn.ReLU())
self.head4 = nn.Sequential(nn.Conv2d(mla_channels, mlahead_channels, 3, padding=1, bias=False),
build_norm_layer(norm_cfg, mlahead_channels)[1], nn.ReLU(),
nn.Conv2d(mlahead_channels, mlahead_channels, 3, padding=1, bias=False),
build_norm_layer(norm_cfg, mlahead_channels)[1], nn.ReLU())
self.head5 = nn.Sequential(nn.Conv2d(mla_channels, mlahead_channels, 3, padding=1, bias=False),
build_norm_layer(norm_cfg, mlahead_channels)[1], nn.ReLU(),
nn.Conv2d(mlahead_channels, mlahead_channels, 3, padding=1, bias=False),
build_norm_layer(norm_cfg, mlahead_channels)[1], nn.ReLU())
def forward(self, mla_p2, mla_p3, mla_p4, mla_p5):
# head2 = self.head2(mla_p2)
head2 = F.interpolate(self.head2(mla_p2), 4*mla_p2.shape[-1], mode='bilinear', align_corners=True)
head3 = F.interpolate(self.head3(mla_p3), 4*mla_p3.shape[-1], mode='bilinear', align_corners=True)
head4 = F.interpolate(self.head4(mla_p4), 4*mla_p4.shape[-1], mode='bilinear', align_corners=True)
head5 = F.interpolate(self.head5(mla_p5), 4*mla_p5.shape[-1], mode='bilinear', align_corners=True)
return torch.cat([head2, head3, head4, head5], dim=1)
@HEADS.register_module()
class VIT_MLAHead(BaseDecodeHead):
""" Vision Transformer with support for patch or hybrid CNN input stage
"""
def __init__(self, img_size=768, mla_channels=256, mlahead_channels=128,
norm_layer=nn.BatchNorm2d, norm_cfg=None, **kwargs):
super(VIT_MLAHead, self).__init__(**kwargs)
self.img_size = img_size
self.norm_cfg = norm_cfg
self.mla_channels = mla_channels
self.BatchNorm = norm_layer
self.mlahead_channels = mlahead_channels
self.mlahead = MLAHead(mla_channels=self.mla_channels, mlahead_channels=self.mlahead_channels, norm_cfg=self.norm_cfg)
self.cls = nn.Conv2d(4 * self.mlahead_channels, self.num_classes, 3, padding=1)
def forward(self, inputs):
x = self.mlahead(inputs[0], inputs[1], inputs[2], inputs[3])
x = self.cls(x)
x = F.interpolate(x, size=self.img_size, mode='bilinear', align_corners=self.align_corners)
return x
|