|
|
|
|
|
|
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
from timm.layers import DropPath, trunc_normal_ |
|
|
|
|
|
|
|
|
class GroupNorm(nn.GroupNorm): |
|
|
""" |
|
|
Group Normalization with 1 group. |
|
|
Input: tensor in shape [B, C, H, W] |
|
|
""" |
|
|
def __init__(self, num_channels, **kwargs): |
|
|
super().__init__(1, num_channels, **kwargs) |
|
|
|
|
|
|
|
|
class Pooling(nn.Module): |
|
|
""" |
|
|
Implementation of pooling for PoolFormer |
|
|
--pool_size: pooling size |
|
|
""" |
|
|
def __init__(self, pool_size=3): |
|
|
super().__init__() |
|
|
self.pool = nn.AvgPool2d( |
|
|
pool_size, stride=1, padding=pool_size//2, count_include_pad=False) |
|
|
|
|
|
def forward(self, x): |
|
|
return self.pool(x) - x |
|
|
|
|
|
|
|
|
class Mlp(nn.Module): |
|
|
""" |
|
|
Implementation of MLP with 1*1 convolutions. |
|
|
Input: tensor with shape [B, C, H, W] |
|
|
""" |
|
|
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.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.Conv2d): |
|
|
trunc_normal_(m.weight, std=.02) |
|
|
if m.bias is not None: |
|
|
nn.init.constant_(m.bias, 0) |
|
|
|
|
|
def forward(self, x): |
|
|
x = self.fc1(x) |
|
|
x = self.act(x) |
|
|
x = self.drop(x) |
|
|
x = self.fc2(x) |
|
|
x = self.drop(x) |
|
|
return x |
|
|
|
|
|
|
|
|
class PoolFormerBlock(nn.Module): |
|
|
""" |
|
|
Implementation of one PoolFormer block. |
|
|
--dim: embedding dim |
|
|
--pool_size: pooling size |
|
|
--mlp_ratio: mlp expansion ratio |
|
|
--act_layer: activation |
|
|
--norm_layer: normalization |
|
|
--drop: dropout rate |
|
|
--drop path: Stochastic Depth, |
|
|
refer to https://arxiv.org/abs/1603.09382 |
|
|
--init_value: LayerScale, |
|
|
refer to https://arxiv.org/abs/2103.17239 |
|
|
""" |
|
|
def __init__(self, dim, pool_size=3, mlp_ratio=4., drop=0., drop_path=0., |
|
|
init_value=1e-5, act_layer=nn.GELU, norm_layer=GroupNorm): |
|
|
super().__init__() |
|
|
|
|
|
self.norm1 = norm_layer(dim) |
|
|
self.token_mixer = Pooling(pool_size=pool_size) |
|
|
self.norm2 = norm_layer(dim) |
|
|
mlp_hidden_dim = int(dim * mlp_ratio) |
|
|
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, |
|
|
act_layer=act_layer, drop=drop) |
|
|
|
|
|
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
|
|
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.token_mixer(self.norm1(x))) |
|
|
x = x + self.drop_path( |
|
|
self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(self.norm2(x))) |
|
|
return x |
|
|
|
