mobilenet V3模型代码

以下是使用PyTorch实现的MobileNet V3模型代码示例：

import torch
import torch.nn as nn

class Hswish(nn.Module):
    def __init__(self, inplace=True):
        super(Hswish, self).__init__()
        self.inplace = inplace

    def forward(self, x):
        if self.inplace:
            x *= torch.clamp(x + 3, 0, 6) / 6
        else:
            x = x * (torch.nn.functional.relu6(x + 3, inplace=True) / 6)
        return x

class Hsigmoid(nn.Module):
    def __init__(self, inplace=True):
        super(Hsigmoid, self).__init__()
        self.inplace = inplace

    def forward(self, x):
        if self.inplace:
            x = torch.clamp(x + 3, 0, 6) / 6
        else:
            x = torch.nn.functional.relu6(x + 3, inplace=True) / 6
        return x

class SeModule(nn.Module):
    def __init__(self, in_size, reduction=4):
        super(SeModule, self).__init__()
        self.se = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_size, in_size // reduction, kernel_size=1, stride=1, padding=0, bias=True),
            nn.Hardswish(inplace=True),
            nn.Conv2d(in_size // reduction, in_size, kernel_size=1, stride=1, padding=0, bias=True),
            Hsigmoid(inplace=True)
        )

    def forward(self, x):
        return x * self.se(x)

class ConvBNActivation(nn.Sequential):
    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, norm_layer=None, activation_layer=None):
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if activation_layer is None:
            activation_layer = nn.ReLU(inplace=True)
        padding = (kernel_size - 1) // 2
        super(ConvBNActivation, self).__init__(
            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
            norm_layer(out_planes),
            activation_layer(inplace=True)
        )

class SqueezeExcite(nn.Module):
    def __init__(self, in_channels, se_ratio=0.25):
        super().__init__()
        self.se_ratio = se_ratio
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        num_squeezed_channels = max(1, int(in_channels * se_ratio))
        self.fc1 = nn.Conv2d(in_channels, num_squeezed_channels, kernel_size=1)
        self.relu = nn.ReLU(inplace=True)
        self.fc2 = nn.Conv2d(num_squeezed_channels, in_channels, kernel_size=1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        out = self.avgpool(x)
        out = self.fc1(out)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.sigmoid(out)
        return out * x

class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio, se_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]
        self.use_res_connect = self.stride == 1 and inp == oup

        layers = []
        # expand
        expand_inp = int(round(inp * expand_ratio))
        if expand_ratio != 1:
            layers.append(ConvBNActivation(inp, expand_inp, kernel_size=1))

        # depthwise
        layers.append(ConvBNActivation(expand_inp, expand_inp, stride=stride, groups=expand_inp, norm_layer=nn.BatchNorm2d, activation_layer=Hswish()))

        # squeeze-and-excite
        layers.append(SqueezeExcite(expand_inp, se_ratio))

        # project
        layers.append(nn.Conv2d(expand_inp, oup, kernel_size=1, bias=False))
        layers.append(nn.BatchNorm2d(oup))

        self.conv = nn.Sequential(*layers)

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)

class MobileNetV3(nn.Module):
    def __init__(self, num_classes=1000, mode='large', width_mult=1.0):
        super(MobileNetV3, self).__init__()
        self.num_classes = num_classes
        if mode == 'large':
            layers = [
                # stage 1
                ConvBNActivation(3, int(16 * width_mult), kernel_size=3, stride=2, norm_layer=nn.BatchNorm2d, activation_layer=Hswish()),

                # stage 2
                InvertedResidual(int(16 * width_mult), int(16 * width_mult), stride=1, expand_ratio=1, se_ratio=0.25),
                InvertedResidual(int(16 * width_mult), int(24 * width_mult), stride=2, expand_ratio=4, se_ratio=0.25),
                InvertedResidual(int(24 * width_mult), int(24 * width_mult), stride=1, expand_ratio=3, se_ratio=0.25),

                # stage 3
                InvertedResidual(int(24 * width_mult), int(40 * width_mult), stride=2, expand_ratio=3, se_ratio=0.25),
                InvertedResidual(int(40 * width_mult), int(40 * width_mult), stride=1, expand_ratio=3, se_ratio=0.25),
                InvertedResidual(int(40 * width_mult), int(40 * width_mult), stride=1, expand_ratio=3, se_ratio=0.25),

                # stage 4
                InvertedResidual(int(40 * width_mult), int(80 * width_mult), stride=2, expand_ratio=6, se_ratio=0.25),
                InvertedResidual(int(80 * width_mult), int(80 * width_mult), stride=1, expand_ratio=2.5, se_ratio=0.25),
                InvertedResidual(int(80 * width_mult), int(80 * width_mult), stride=1, expand_ratio=2.3, se_ratio=0.25),
                InvertedResidual(int(80 * width_mult), int(80 * width_mult), stride=1, expand_ratio=2.3, se_ratio=0.25),

                # stage 5
                InvertedResidual(int(80 * width_mult), int(112 * width_mult), stride=1, expand_ratio=6, se_ratio=0.25),
                InvertedResidual(int(112 * width_mult), int(112 * width_mult), stride=1, expand_ratio=6, se_ratio=0.25),

                # stage 6
                InvertedResidual(int(112 * width_mult), int(160 * width_mult), stride=2, expand_ratio=6, se_ratio=0.25),
                InvertedResidual(int(160 * width_mult), int(160 * width_mult), stride=1, expand_ratio=6, se_ratio=0.25),
                InvertedResidual(int(160 * width_mult), int(160 * width_mult), stride=1, expand_ratio=6, se_ratio=0.25),

                # stage 7
                ConvBNActivation(int(160 * width_mult), int(960 * width_mult), kernel_size=1, norm_layer=nn.BatchNorm2d, activation_layer=Hswish())
            ]
        elif mode == 'small':
            layers = [
                # stage 1
                ConvBNActivation(3, int(8 * width_mult), kernel_size=3, stride=2, norm_layer=nn.BatchNorm2d, activation_layer=Hswish()),

                # stage 2
                InvertedResidual(int(8 * width_mult), int(16 * width_mult), stride=2, expand_ratio=3, se_ratio=0.25),
                InvertedResidual(int(16 * width_mult), int(16 * width_mult), stride=1, expand_ratio=3, se_ratio=0.25),

                # stage 3
                InvertedResidual(int(16 * width_mult), int(24 * width_mult), stride=2, expand_ratio=3, se_ratio=0.25),
                InvertedResidual(int(24 * width_mult), int(24 * width_mult), stride=1, expand_ratio=3, se_ratio=0.25),

                # stage 4
                InvertedResidual(int(24 * width_mult), int(40 * width_mult), stride=2, expand_ratio=6, se_ratio=0.25),
                InvertedResidual(int(40 * width_mult), int(40 * width_mult), stride=1, expand_ratio=6, se_ratio=0.25),
                InvertedResidual(int(40 * width_mult), int(40 * width_mult), stride=1, expand_ratio=6, se_ratio=0.25),

                # stage 5
                InvertedResidual(int(40 * width_mult), int(48 * width_mult), stride=1, expand_ratio=6, se_ratio=0.25),
                InvertedResidual(int(48 * width_mult), int(48 * width_mult), stride=1, expand_ratio=6, se_ratio=0.25),

                # stage 6
                InvertedResidual(int(48 * width_mult), int(96 * width_mult), stride=2, expand_ratio=6, se_ratio=0.25),
                InvertedResidual(int(96 * width_mult), int(96 * width_mult), stride=1, expand_ratio=6, se_ratio=0.25),

                # stage 7
                InvertedResidual(int(96 * width_mult), int(576 * width_mult), kernel_size=1, stride=1, expand_ratio=6, se_ratio=0.25),
                ConvBNActivation(int(576 * width_mult), int(1024 * width_mult), kernel_size=1, norm_layer=nn.BatchNorm2d, activation_layer=Hswish())
            ]
        else:
            raise ValueError('Unsupported MobileNetV3 model mode: ', mode)

        # build the network
        self.features = nn.Sequential(*layers)

        # build the head
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Sequential(
            nn.Linear(int(1024 * width_mult), int(1280 * width_mult)),
            Hswish(inplace=True),
            nn.Dropout(p=0.2, inplace=True) if self.num_classes > 0 else nn.Identity(),
            nn.Linear(int(1280 * width_mult), self.num_classes),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x