diff --git a/models/__init__.py b/models/__init__.py index 61bdf57..89e83b1 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -8,6 +8,7 @@ from .mobilenet import * from .mobilenetv2 import * from .mobilenetv3 import * +from .ghostnet import * from .shufflenet import * from .shufflenetv2 import * from .mnasnet import * diff --git a/models/ghostnet.py b/models/ghostnet.py new file mode 100644 index 0000000..139dc72 --- /dev/null +++ b/models/ghostnet.py @@ -0,0 +1,225 @@ +"""https://github.com/huawei-noah/CV-Backbones/blob/master/ghostnet_pytorch/ghostnet.py""" +import os +import math +import torch +import torch.nn as nn +from .core import blocks, export, make_divisible +from typing import Any, List + + +class GhostModule(nn.Module): + def __init__( + self, + inp, + oup, + kernel_size=1, + ratio=2, + dw_size=3, + stride=1, + relu=True + ): + super().__init__() + super(GhostModule, self).__init__() + self.oup = oup + init_channels = math.ceil(oup / ratio) + new_channels = init_channels * (ratio - 1) + + self.primary_conv = nn.Sequential( + nn.Conv2d(inp, init_channels, kernel_size, stride, kernel_size//2, bias=False), + nn.BatchNorm2d(init_channels), + nn.ReLU(inplace=True) if relu else nn.Sequential(), + ) + + self.cheap_operation = nn.Sequential( + nn.Conv2d(init_channels, new_channels, dw_size, 1, dw_size//2, groups=init_channels, bias=False), + nn.BatchNorm2d(new_channels), + nn.ReLU(inplace=True) if relu else nn.Sequential(), + ) + + def forward(self, x): + x1 = self.primary_conv(x) + x2 = self.cheap_operation(x1) + out = torch.cat([x1,x2], dim=1) + return out[:,:self.oup,:,:] + + +class GhostBottleneck(nn.Module): + """ Ghost bottleneck w/ optional SE""" + + def __init__( + self, + in_chs, + mid_chs, + out_chs, + dw_kernel_size=3, + stride=1, + act_layer=nn.ReLU, + se_ratio=0. + ): + super(GhostBottleneck, self).__init__() + has_se = se_ratio is not None and se_ratio > 0. + self.stride = stride + + # Point-wise expansion + self.ghost1 = GhostModule(in_chs, mid_chs, relu=True) + + # 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) + + # Squeeze-and-excitation + if has_se: + self.se = blocks.SEBlock(mid_chs, ratio=se_ratio) + else: + self.se = None + + # Point-wise linear projection + self.ghost2 = GhostModule(mid_chs, out_chs, relu=False) + + # shortcut + if (in_chs == out_chs and self.stride == 1): + self.shortcut = nn.Sequential() + else: + self.shortcut = nn.Sequential( + nn.Conv2d(in_chs, in_chs, dw_kernel_size, stride=stride, + padding=(dw_kernel_size-1)//2, groups=in_chs, bias=False), + nn.BatchNorm2d(in_chs), + nn.Conv2d(in_chs, out_chs, 1, stride=1, padding=0, bias=False), + nn.BatchNorm2d(out_chs), + ) + + def forward(self, x): + residual = x + + # 1st ghost bottleneck + x = self.ghost1(x) + + # Depth-wise convolution + if self.stride > 1: + 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(residual) + return x + + +@export +class GhostNet(nn.Module): + def __init__( + self, + in_channels: int = 3, + num_classes: int = 1000, + multiplier: float = 1.0, + dropout_rate: float = 0.2, + cfgs: List[list] = [], + thumbnail: bool = False + ) -> None: + super().__init__() + + FRONT_S = 1 if thumbnail else 2 + + inp = make_divisible(16 * multiplier, 4) + _features = [blocks.Conv2dBlock(in_channels, inp, stride=FRONT_S)] + + for cfg in cfgs: + _layers = [] + for k, t, c, se_ratio, s in cfg: + oup = make_divisible(c * multiplier, 4) + _layers.append(GhostBottleneck(inp, make_divisible(t * multiplier, 4), oup, k, s, se_ratio=se_ratio)) + inp = oup + + _features.append(blocks.Stage(*_layers)) + + oup = make_divisible(t * multiplier, 4) + _features.append(blocks.Conv2d1x1Block(inp, oup)) + + self.features = nn.Sequential(*_features) + + self.avg = nn.AdaptiveAvgPool2d((1, 1)) + self.classifier = nn.Sequential( + nn.Linear(oup, 1280), + nn.ReLU(inplace=True), + nn.Dropout(dropout_rate, inplace=True), + nn.Linear(1280, num_classes) + ) + + def forward(self, x): + x = self.features(x) + x = self.avg(x) + x = torch.flatten(x, 1) + x = self.classifier(x) + return x + + +def _ghostnet( + multiplier: float = 1.0, + pretrained: bool = False, + pth: str = None, + progress: bool = True, + **kwargs: Any +): + cfgs = [ + # k, t, c, SE, s + # stage1 + [[3, 16, 16, 0, 1]], + # stage2 + [[3, 48, 24, 0, 2]], + [[3, 72, 24, 0, 1]], + # stage3 + [[5, 72, 40, 0.25, 2]], + [[5, 120, 40, 0.25, 1]], + # stage4 + [[3, 240, 80, 0, 2]], + [[3, 200, 80, 0, 1], + [3, 184, 80, 0, 1], + [3, 184, 80, 0, 1], + [3, 480, 112, 0.25, 1], + [3, 672, 112, 0.25, 1] + ], + # stage5 + [[5, 672, 160, 0.25, 2]], + [[5, 960, 160, 0, 1], + [5, 960, 160, 0.25, 1], + [5, 960, 160, 0, 1], + [5, 960, 160, 0.25, 1] + ] + ] + + model = GhostNet(multiplier=multiplier, cfgs=cfgs, **kwargs) + + if pretrained: + if pth is not None: + state_dict = torch.load(os.path.expanduser(pth)) + else: + assert 'url' in kwargs and kwargs['url'] != '', 'Invalid URL.' + state_dict = torch.hub.load_state_dict_from_url( + kwargs['url'], + progress=progress + ) + model.load_state_dict(state_dict) + return model + + +@export +def ghostnet_x0_5(pretrained: bool = False, pth: str = None, progress: bool = True, **kwargs): + return _ghostnet(0.5, pretrained, pth, progress, **kwargs) + + +@export +def ghostnet_x1_0(pretrained: bool = False, pth: str = None, progress: bool = True, **kwargs): + return _ghostnet(1.0, pretrained, pth, progress, **kwargs) + + +@export +def ghostnet_x1_3(pretrained: bool = False, pth: str = None, progress: bool = True, **kwargs): + return _ghostnet(1.3, pretrained, pth, progress, **kwargs)