""" An implementation of GhostNet Model as defined in: GhostNet: More Features from Cheap Operations. https://arxiv.org/abs/1911.11907 The train script of the model is similar to that of MobileNetV3 Original model: https://github.com/huawei-noah/CV-backbones/tree/master/ghostnet_pytorch """ import math from functools import partial import torch import torch.nn as nn import torch.nn.functional as F from timm.layers import drop_path, SqueezeExcite from timm.models.layers import CondConv2d, hard_sigmoid, DropPath __all__ = ['RepNCSPELAN4_GhostModule'] _SE_LAYER = partial(SqueezeExcite, gate_fn=hard_sigmoid, divisor=4) class AdaptiveSE(nn.Module): def __init__(self, channels, reduction=4): super(AdaptiveSE, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(inplace=True), nn.Linear(channels // reduction, channels), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() y = self.avg_pool(x).view(b, c) y = self.fc(y).view(b, c, 1, 1) return x * y.expand_as(x) class DynamicActivation(nn.Module): def __init__(self): super(DynamicActivation, self).__init__() self.act_fns = nn.ModuleDict({ 'relu': nn.ReLU(inplace=True), 'leaky_relu': nn.LeakyReLU(0.1, inplace=True), 'silu': nn.SiLU(inplace=True) }) self.selected_act = 'relu' def forward(self, x): return self.act_fns[self.selected_act](x) def set_activation(self, act_name): if act_name in self.act_fns: self.selected_act = act_name else: raise ValueError("Activation function not supported") class DynamicConv(nn.Module): """ Dynamic Conv layer """ def __init__(self, in_features, out_features, kernel_size=1, stride=1, padding='', dilation=1, groups=1, bias=False, num_experts=4): super().__init__() self.routing = nn.Linear(in_features, num_experts) self.cond_conv = CondConv2d(in_features, out_features, kernel_size, stride, padding, dilation, groups, bias, num_experts) def forward(self, x): pooled_inputs = F.adaptive_avg_pool2d(x, 1).flatten(1) # CondConv routing routing_weights = torch.sigmoid(self.routing(pooled_inputs)) x = self.cond_conv(x, routing_weights) return x class ConvBnAct(nn.Module): """ Conv + Norm Layer + Activation w/ optional skip connection """ def __init__( self, in_chs, out_chs, kernel_size, stride=1, dilation=1, pad_type='', skip=False, act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, drop_path_rate=0., num_experts=4): super(ConvBnAct, self).__init__() self.has_residual = skip and stride == 1 and in_chs == out_chs self.drop_path_rate = drop_path_rate # self.conv = create_conv2d(in_chs, out_chs, kernel_size, stride=stride, dilation=dilation, padding=pad_type) self.conv = DynamicConv(in_chs, out_chs, kernel_size, stride, dilation=dilation, padding=pad_type, num_experts=num_experts) self.bn1 = norm_layer(out_chs) self.act1 = act_layer() def feature_info(self, location): if location == 'expansion': # output of conv after act, same as block coutput info = dict(module='act1', hook_type='forward', num_chs=self.conv.out_channels) else: # location == 'bottleneck', block output info = dict(module='', hook_type='', num_chs=self.conv.out_channels) return info def forward(self, x): shortcut = x x = self.conv(x) x = self.bn1(x) x = self.act1(x) if self.has_residual: if self.drop_path_rate > 0.: x = drop_path(x, self.drop_path_rate, self.training) x += shortcut return x class GhostModule(nn.Module): def __init__(self, inp, oup, kernel_size=1, ratio=2, dw_size=3, stride=1, act_layer=nn.ReLU, num_experts=4): super(GhostModule, self).__init__() self.oup = oup init_channels = math.ceil(oup / ratio) new_channels = init_channels * (ratio - 1) self.primary_conv = nn.Sequential( DynamicConv(inp, init_channels, kernel_size, stride, kernel_size // 2, bias=False, num_experts=num_experts), nn.BatchNorm2d(init_channels), DynamicActivation() ) self.cheap_operation = nn.Sequential( DynamicConv(init_channels, new_channels, dw_size, 1, dw_size // 2, groups=init_channels, bias=False, num_experts=num_experts), nn.BatchNorm2d(new_channels), DynamicActivation() ) self.se_module = AdaptiveSE(init_channels + new_channels) def forward(self, x): x1 = self.primary_conv(x) x2 = self.cheap_operation(x1) out = torch.cat([x1, x2], dim=1) out = self.se_module(out) return out[:, :self.oup, :, :] class GhostBottleneck(nn.Module): """ Ghost bottleneck w/ optional SE""" def __init__(self, in_chs, out_chs, dw_kernel_size=3, stride=1, act_layer=nn.ReLU, se_ratio=0., drop_path=0., num_experts=4): super(GhostBottleneck, self).__init__() has_se = se_ratio is not None and se_ratio > 0. self.stride = stride mid_chs = in_chs * 2 # Point-wise expansion self.ghost1 = GhostModule(in_chs, mid_chs, act_layer=act_layer, num_experts=num_experts) # Depth-wise convolution if self.stride > 1: self.conv_dw = nn.Conv2d( mid_chs, mid_chs, dw_kernel_size, stride=stride, padding=(dw_kernel_size - 1) // 2, groups=mid_chs, bias=False) self.bn_dw = nn.BatchNorm2d(mid_chs) else: self.conv_dw = None self.bn_dw = None # Squeeze-and-excitation self.se = _SE_LAYER(mid_chs, se_ratio=se_ratio, act_layer=act_layer if act_layer is not nn.GELU else nn.ReLU) if has_se else None # Point-wise linear projection self.ghost2 = GhostModule(mid_chs, out_chs, act_layer=None, num_experts=num_experts) # shortcut if in_chs == out_chs and self.stride == 1: self.shortcut = nn.Sequential() else: self.shortcut = nn.Sequential( DynamicConv( in_chs, in_chs, dw_kernel_size, stride=stride, padding=(dw_kernel_size - 1) // 2, groups=in_chs, bias=False, num_experts=num_experts), nn.BatchNorm2d(in_chs), DynamicConv(in_chs, out_chs, 1, stride=1, padding=0, bias=False, num_experts=num_experts), nn.BatchNorm2d(out_chs), ) self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() def forward(self, x): shortcut = x # 1st ghost bottleneck x = self.ghost1(x) # Depth-wise convolution if self.conv_dw is not None: x = self.conv_dw(x) x = self.bn_dw(x) # Squeeze-and-excitation if self.se is not None: x = self.se(x) # 2nd ghost bottleneck x = self.ghost2(x) x = self.shortcut(shortcut) + self.drop_path(x) return x def autopad(k, p=None, d=1): # kernel, padding, dilation """Pad to 'same' shape outputs.""" if d > 1: k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size if p is None: p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad return p class Conv(nn.Module): """Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation).""" default_act = nn.SiLU() # default activation def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True): """Initialize Conv layer with given arguments including activation.""" super().__init__() self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False) self.bn = nn.BatchNorm2d(c2) self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity() def forward(self, x): """Apply convolution, batch normalization and activation to input tensor.""" return self.act(self.bn(self.conv(x))) def forward_fuse(self, x): """Perform transposed convolution of 2D data.""" return self.act(self.conv(x)) class RepNCSP(nn.Module): # CSP Bottleneck with 3 convolutions def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion super().__init__() c_ = int(c2 * e) # hidden channels self.cv1 = Conv(c1, c_, 1, 1) self.cv2 = Conv(c1, c_, 1, 1) self.cv3 = Conv(2 * c_, c2, 1) # optional act=FReLU(c2) self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n))) def forward(self, x): return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1)) class RepNCSPELAN4_GhostModule(nn.Module): # csp-elan def __init__(self, c1, c2, c3, c4, c5=1): # ch_in, ch_out, number, shortcut, groups, expansion super().__init__() self.c = c3//2 self.cv1 = Conv(c1, c3, 1, 1) self.cv2 = nn.Sequential(RepNCSP(c3//2, c4, c5), Conv(c4, c4, 3, 1)) self.cv3 = nn.Sequential(RepNCSP(c4, c4, c5), Conv(c4, c4, 3, 1)) self.cv4 = Conv(c3+(2*c4), c2, 1, 1) def forward(self, x): y = list(self.cv1(x).chunk(2, 1)) y.extend((m(y[-1])) for m in [self.cv2, self.cv3]) return self.cv4(torch.cat(y, 1)) def forward_split(self, x): y = list(self.cv1(x).split((self.c, self.c), 1)) y.extend(m(y[-1]) for m in [self.cv2, self.cv3]) return self.cv4(torch.cat(y, 1)) if __name__ == "__main__": # Generating Sample image image_size = (1, 64, 224, 224) image = torch.rand(*image_size) # Model model = RepNCSPELAN4_GhostModule(64, 64) out = model(image) print(out.size())