|
|
|
|
|
|
|
|
|
|
|
import math |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
from timm.layers import DropPath, trunc_normal_ |
|
|
|
|
|
|
|
|
class DWConv(nn.Module): |
|
|
def __init__(self, dim=768): |
|
|
super(DWConv, self).__init__() |
|
|
self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim) |
|
|
|
|
|
def forward(self, x): |
|
|
x = self.dwconv(x) |
|
|
return x |
|
|
|
|
|
|
|
|
class MixMlp(nn.Module): |
|
|
def __init__(self, |
|
|
in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): |
|
|
super().__init__() |
|
|
out_features = out_features or in_features |
|
|
hidden_features = hidden_features or in_features |
|
|
self.fc1 = nn.Conv2d(in_features, hidden_features, 1) |
|
|
self.dwconv = DWConv(hidden_features) |
|
|
self.act = act_layer() |
|
|
self.fc2 = nn.Conv2d(hidden_features, out_features, 1) |
|
|
self.drop = nn.Dropout(drop) |
|
|
self.apply(self._init_weights) |
|
|
|
|
|
def _init_weights(self, m): |
|
|
if isinstance(m, nn.Linear): |
|
|
trunc_normal_(m.weight, std=.02) |
|
|
if isinstance(m, nn.Linear) and m.bias is not None: |
|
|
nn.init.constant_(m.bias, 0) |
|
|
elif isinstance(m, nn.LayerNorm): |
|
|
nn.init.constant_(m.bias, 0) |
|
|
nn.init.constant_(m.weight, 1.0) |
|
|
elif isinstance(m, nn.Conv2d): |
|
|
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels |
|
|
fan_out //= m.groups |
|
|
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) |
|
|
if m.bias is not None: |
|
|
m.bias.data.zero_() |
|
|
|
|
|
def forward(self, x): |
|
|
x = self.fc1(x) |
|
|
x = self.dwconv(x) |
|
|
x = self.act(x) |
|
|
x = self.drop(x) |
|
|
x = self.fc2(x) |
|
|
x = self.drop(x) |
|
|
return x |
|
|
|
|
|
|
|
|
class LKA(nn.Module): |
|
|
def __init__(self, dim): |
|
|
super().__init__() |
|
|
self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim) |
|
|
self.conv_spatial = nn.Conv2d( |
|
|
dim, dim, 7, stride=1, padding=9, groups=dim, dilation=3) |
|
|
self.conv1 = nn.Conv2d(dim, dim, 1) |
|
|
|
|
|
|
|
|
def forward(self, x): |
|
|
u = x.clone() |
|
|
attn = self.conv0(x) |
|
|
attn = self.conv_spatial(attn) |
|
|
attn = self.conv1(attn) |
|
|
|
|
|
return u * attn |
|
|
|
|
|
|
|
|
class Attention(nn.Module): |
|
|
def __init__(self, d_model, attn_shortcut=True): |
|
|
super().__init__() |
|
|
|
|
|
self.proj_1 = nn.Conv2d(d_model, d_model, 1) |
|
|
self.activation = nn.GELU() |
|
|
self.spatial_gating_unit = LKA(d_model) |
|
|
self.proj_2 = nn.Conv2d(d_model, d_model, 1) |
|
|
self.attn_shortcut = attn_shortcut |
|
|
|
|
|
def forward(self, x): |
|
|
if self.attn_shortcut: |
|
|
shortcut = x.clone() |
|
|
x = self.proj_1(x) |
|
|
x = self.activation(x) |
|
|
x = self.spatial_gating_unit(x) |
|
|
x = self.proj_2(x) |
|
|
if self.attn_shortcut: |
|
|
x = x + shortcut |
|
|
return x |
|
|
|
|
|
|
|
|
class VANBlock(nn.Module): |
|
|
def __init__(self, dim, mlp_ratio=4., drop=0.,drop_path=0., init_value=1e-2, act_layer=nn.GELU, attn_shortcut=True): |
|
|
super().__init__() |
|
|
self.norm1 = nn.BatchNorm2d(dim) |
|
|
self.attn = Attention(dim, attn_shortcut=attn_shortcut) |
|
|
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
|
|
|
|
|
self.norm2 = nn.BatchNorm2d(dim) |
|
|
mlp_hidden_dim = int(dim * mlp_ratio) |
|
|
self.mlp = MixMlp( |
|
|
in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
|
|
|
|
|
self.layer_scale_1 = nn.Parameter(init_value * torch.ones((dim)), requires_grad=True) |
|
|
self.layer_scale_2 = nn.Parameter(init_value * torch.ones((dim)), requires_grad=True) |
|
|
|
|
|
def forward(self, x): |
|
|
x = x + self.drop_path( |
|
|
self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * self.attn(self.norm1(x))) |
|
|
x = x + self.drop_path( |
|
|
self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(self.norm2(x))) |
|
|
return x |
|
|
|
