图像分类MobileNet系列源代码:v1-v3

from typing import Callable, List, Optional import torch from torch import nn, Tensor from torch.nn import functional as F from functools import partial def _make_divisible(ch, divisor=8, min_ch=None): """ This function is taken from the original tf repo. It ensures that all layers have a channel number that is divisible by 8 It can be seen here: https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py """ if min_ch is None: min_ch = divisor new_ch = max(min_ch, int(ch + divisor / 2) // divisor * divisor) # Make sure that round down does not go down by more than 10%. if new_ch < 0.9 * ch: new_ch += divisor return new_ch class ConvBNActivation(nn.Sequential): def __init__(self, in_planes: int, out_planes: int, kernel_size: int = 3, stride: int = 1, groups: int = 1, norm_layer: Optional[Callable[..., nn.Module]] = None, activation_layer: Optional[Callable[..., nn.Module]] = None): padding = (kernel_size - 1) // 2 if norm_layer is None: norm_layer = nn.BatchNorm2d if activation_layer is None: activation_layer = nn.ReLU6 super(ConvBNActivation, self).__init__(nn.Conv2d(in_channels=in_planes, out_channels=out_planes, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, bias=False), norm_layer(out_planes), activation_layer(inplace=True)) class SqueezeExcitation(nn.Module): def __init__(self, input_c: int, squeeze_factor: int = 4): super(SqueezeExcitation, self).__init__() squeeze_c = _make_divisible(input_c // squeeze_factor, 8) self.fc1 = nn.Conv2d(input_c, squeeze_c, 1) self.fc2 = nn.Conv2d(squeeze_c, input_c, 1) def forward(self, x: Tensor) -> Tensor: scale = F.adaptive_avg_pool2d(x, output_size=(1, 1)) scale = self.fc1(scale) scale = F.relu(scale, inplace=True) scale = self.fc2(scale) scale = F.hardsigmoid(scale, inplace=True) return scale * x class InvertedResidualConfig: def __init__(self, input_c: int, kernel: int, expanded_c: int, out_c: int, use_se: bool, activation: str, stride: int, width_multi: float): self.input_c = self.adjust_channels(input_c, width_multi) self.kernel = kernel self.expanded_c = self.adjust_channels(expanded_c, width_multi) self.out_c = self.adjust_channels(out_c, width_multi) self.use_se = use_se self.use_hs = activation == "HS" # whether using h-swish activation self.stride = stride @staticmethod def adjust_channels(channels: int, width_multi: float): return _make_divisible(channels * width_multi, 8) class InvertedResidual(nn.Module): def __init__(self, cnf: InvertedResidualConfig, norm_layer: Callable[..., nn.Module]): super(InvertedResidual, self).__init__() if cnf.stride not in [1, 2]: raise ValueError("illegal stride value.") self.use_res_connect = (cnf.stride == 1 and cnf.input_c == cnf.out_c) layers: List[nn.Module] = [] activation_layer = nn.Hardswish if cnf.use_hs else nn.ReLU # expand if cnf.expanded_c != cnf.input_c: layers.append(ConvBNActivation(cnf.input_c, cnf.expanded_c, kernel_size=1, norm_layer=norm_layer, activation_layer=activation_layer)) # depthwise layers.append(ConvBNActivation(cnf.expanded_c, cnf.expanded_c, kernel_size=cnf.kernel, stride=cnf.stride, groups=cnf.expanded_c, norm_layer=norm_layer, activation_layer=activation_layer)) if cnf.use_se: layers.append(SqueezeExcitation(cnf.expanded_c)) # project layers.append(ConvBNActivation(cnf.expanded_c, cnf.out_c, kernel_size=1, norm_layer=norm_layer, activation_layer=nn.Identity)) self.block = nn.Sequential(*layers) self.out_channels = cnf.out_c self.is_strided = cnf.stride > 1 def forward(self, x: Tensor) -> Tensor: result = self.block(x) if self.use_res_connect: result += x return result class MobileNetV3(nn.Module): def __init__(self, inverted_residual_setting: List[InvertedResidualConfig], last_channel: int, num_classes: int = 1000, block: Optional[Callable[..., nn.Module]] = None, norm_layer: Optional[Callable[..., nn.Module]] = None): super(MobileNetV3, self).__init__() if not inverted_residual_setting: raise ValueError("The inverted_residual_setting should not be empty.") elif not (isinstance(inverted_residual_setting, List) and all([isinstance(s, InvertedResidualConfig) for s in inverted_residual_setting])): raise TypeError("The inverted_residual_setting should be List[InvertedResidualConfig]") if block is None: block = InvertedResidual if norm_layer is None: norm_layer = partial(nn.BatchNorm2d, eps=0.001, momentum=0.01) layers: List[nn.Module] = [] # building first layer firstconv_output_c = inverted_residual_setting[0].input_c layers.append(ConvBNActivation(3, firstconv_output_c, kernel_size=3, stride=2, norm_layer=norm_layer, activation_layer=nn.Hardswish)) # building inverted residual blocks for cnf in inverted_residual_setting: layers.append(block(cnf, norm_layer)) # building last several layers lastconv_input_c = inverted_residual_setting[-1].out_c lastconv_output_c = 6 * lastconv_input_c layers.append(ConvBNActivation(lastconv_input_c, lastconv_output_c, kernel_size=1, norm_layer=norm_layer, activation_layer=nn.Hardswish)) self.features = nn.Sequential(*layers) self.avgpool = nn.AdaptiveAvgPool2d(1) self.classifier = nn.Sequential(nn.Linear(lastconv_output_c, last_channel), nn.Hardswish(inplace=True), nn.Dropout(p=0.2, inplace=True), nn.Linear(last_channel, num_classes)) # initial weights for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode="fan_out")