mobilenetv2_unet

mobileunet：

https://github.com/YZY-stack/UNet-MobileNet-Pytorch

mobilenetv2：

FootTrackingAR/ImgWrapNet.py at 87c73cc5b64d9a2227792334c479a51696a9d70b · crea9813/FootTrackingAR · GitHub

cpu 320*320需要3s
 

https://github.com/HCMUS-ROBOTICS/ssdf-nncore/blob/532c7b6521d35c8fcb6412c08766e286ad04ef27/nncore/segmentation/models/mobileunet.py

这个还稍微靠谱点：

import time
from collections import OrderedDict

import torch

from torch import nn

def _make_divisible(v, divisor, min_value=None):
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v


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

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

        if norm_layer is None:
            norm_layer = nn.BatchNorm2d

        hidden_dim = int(round(inp * expand_ratio))
        self.use_res_connect = self.stride == 1 and inp == oup

        layers = []
        if expand_ratio != 1:
            # pw
            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer))
        layers.extend([# dw
            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, norm_layer=norm_layer), # pw-linear
            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), norm_layer(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 MobileNetV2(nn.Module):
    def __init__(self, width_mult=1.0, inverted_residual_setting=None, round_nearest=8, block=InvertedResidual, norm_layer=nn.BatchNorm2d):
        super(MobileNetV2, self).__init__()

        input_channel = 32
        last_channel = 240

        if inverted_residual_setting is None:
            inverted_residual_setting = [# t, c, n, s
                [1, 16, 1, 1], [6, 24, 2, 2], [6, 32, 3, 2], [6, 64, 4, 2], [6, 96, 3, 1], [6, 160, 3, 2], [6, 240, 1, 1], ]

        # only check the first element, assuming user knows t,c,n,s are required
        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
            raise ValueError("inverted_residual_setting should be non-empty "
                             "or a 4-element list, got {}".format(inverted_residual_setting))

        # building first layer
        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
        features = [ConvBNReLU(3, input_channel, stride=2, norm_layer=norm_layer)]
        # building inverted residual blocks
        for t, c, n, s in inverted_residual_setting:
            output_channel = _make_divisible(c * width_mult, round_nearest)
            for i in range(n):
                stride = s if i == 0 else 1
                features.append(block(input_channel, output_channel, stride, expand_ratio=t, norm_layer=norm_layer))
                input_channel = output_channel
        # building last several layers
        features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1, norm_layer=norm_layer))
        # make it nn.Sequential
        self.features = nn.Sequential(*features)


        # weight initialization
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.zeros_(m.bias)

    def forward(self, x):
        x = self.features(x)
        return x


class InvertedResidual_a(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual_a, self).__init__()
        self.stride = stride
        assert stride in [1, 2]
        hidden_dim = round(inp * expand_ratio)
        self.use_res_connect = self.stride == 1 and inp == oup
        if expand_ratio == 1:
            self.conv = nn.Sequential(  # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False), nn.BatchNorm2d(hidden_dim), nn.ReLU6(inplace=True),  # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup), )
        else:
            self.conv = nn.Sequential(  # pw
                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False), nn.BatchNorm2d(hidden_dim), nn.ReLU6(inplace=True),  # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False), nn.BatchNorm2d(hidden_dim), nn.ReLU6(inplace=True),  # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup), )

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


class MobileUnet(nn.Module):
    def __init__(self):
        super(MobileUnet, self).__init__()
        mobilenet = MobileNetV2()

        self.down1 = nn.Sequential(*mobilenet.features[0:2])
        self.down2 = nn.Sequential(*mobilenet.features[2:4])
        self.down3 = nn.Sequential(*mobilenet.features[4:7])
        self.down4 = nn.Sequential(*mobilenet.features[7:14])
        self.down5 = nn.Sequential(*mobilenet.features[14:19])

        self.dconv1 = nn.ConvTranspose2d(240, 96, 4, padding=1, stride=2)
        self.invres1 = InvertedResidual_a(192, 96, 1, 6)

        self.dconv2 = nn.ConvTranspose2d(96, 32, 4, padding=1, stride=2)
        self.invres2 = InvertedResidual_a(64, 32, 1, 6)

        self.dconv3 = nn.ConvTranspose2d(32, 24, 4, padding=1, stride=2)
        self.invres3 = InvertedResidual_a(48, 24, 1, 6)

        self.dconv4 = nn.ConvTranspose2d(24, 16, 4, padding=1, stride=2)
        self.invres4 = InvertedResidual_a(32, 16, 1, 6)

        self.dconv5 = nn.ConvTranspose2d(16, 3, 4, padding=1, stride=2)
        self.invres5 = InvertedResidual_a(6, 3, 1, 6)

        self.conv_last = nn.Conv2d(3, 1, 1)

    def _init_weights(self):
        for module in self.modules():
            if isinstance(module, Up):
                module._init_weights()

    def forward(self, x):
        x0 = x
        x1 = x = self.down1(x)
        x2 = x = self.down2(x)
        # print((x2.shape, 'x2'))
        x3 = x = self.down3(x)
        x4 = x = self.down4(x)
        # print((x4.shape, 'x4'))
        x = self.down5(x)
        # print((x5.shape, 'x5'))

        up1 = torch.cat([x4, self.dconv1(x)], dim=1)
        up1 = self.invres1(up1)
        # print((up1.shape, 'up1'))

        up2 = torch.cat([x3, self.dconv2(up1)], dim=1)
        up2 = self.invres2(up2)
        # print((up2.shape, 'up2'))

        up3 = torch.cat([x2, self.dconv3(up2)], dim=1)
        up3 = self.invres3(up3)
        # print((up3.shape, 'up3'))

        up4 = torch.cat([x1, self.dconv4(up3)], dim=1)
        up4 = self.invres4(up4)
        # print((up4.shape, 'up4'))

        up5 = torch.cat([x0, self.dconv5(up4)], dim=1)
        up5 = self.invres5(up5)

        x = self.conv_last(up5)
        return x


class Up(nn.Module):
    def __init__(self, in_channels, in_concat_channels, out_channels):
        super(Up, self).__init__()
        self.in_channels = in_channels
        self.in_concat_channels = in_concat_channels
        self.out_channels = out_channels

        self.up_conv = UpConv(in_channels)
        self.double_conv = DoubleConv(self.in_channels + self.in_concat_channels, self.out_channels)

    def forward(self, x1, x2):
        """
        :param x1: feature from previous layer
        :param x2: feature to concat
        """
        x1 = self.up_conv(x1)
        x = torch.cat([x2, x1], dim=1)
        x = self.double_conv(x)
        return x


class UpConv(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.up_conv = nn.Sequential(OrderedDict([("up", nn.Upsample(scale_factor=2)), ("conv", nn.Conv2d(in_channels, in_channels // 2, kernel_size=3, padding=1),), ]))

    def forward(self, x):
        return self.up_conv(x)


class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True),
                                         nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), )

    def forward(self, x):
        return self.double_conv(x)


if __name__ == "__main__":

    model = MobileUnet().cuda()

    input = torch.randn(1, 3, 448, 448).cuda()

    for i in range(20):
        start = time.time()
        output = model(input)
        print('output.size ', output.size(), time.time() - start)
        del output