pytorch复现MobieNetV2


python 复制代码
import torch
import torch.nn as nn


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 ConvBNReLU(nn.Sequential):
    def __init__(self, in_channel, out_channel, kernel_size=3, stride=1, groups=1):
        padding = (kernel_size - 1) // 2
        super(ConvBNReLU, self).__init__(
            nn.Conv2d(in_channel, out_channel, kernel_size, stride, padding, groups=groups, bias=False),
            nn.BatchNorm2d(out_channel),
            nn.ReLU6(inplace=True)
        )


class InvertedResidual(nn.Module):
    def __init__(self, in_channel, out_channel, stride, expand_ration):
        super(InvertedResidual, self).__init__()
        hidden_channel = in_channel * expand_ration
        self.use_shortcut = (stride==1 and in_channel==out_channel)

        layers = []

        if expand_ration != 1:
            layers.append(ConvBNReLU(in_channel, hidden_channel, kernel_size=1))

        #extent 添加多个元素
        layers.extend([
            #group=1 就是普通卷积, group=hidden_channel 就是depthwise卷积
            ConvBNReLU(hidden_channel, hidden_channel, stride=stride, groups=hidden_channel),
            nn.Conv2d(hidden_channel, out_channel, kernel_size=1, bias=False),
            nn.BatchNorm2d(out_channel)
        ])

        self.conv = nn.Sequential(*layers)


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


class MobileNetV2(nn.Module):
    def __init__(self, num_classes=1000, alpha=1.0, round_nearest=8):
        super(MobileNetV2, self).__init__()
        block = InvertedResidual
        #使通道数32*alpha 变成最接近 round_nearest 的整倍数
        input_channel = _make_divisible(32 * alpha, round_nearest)
        last_channel = _make_divisible(1280 * alpha, round_nearest)

        '''
            t:扩展因子
            c:输出特征矩阵深度channel
            n:bottleneck的重复次数
            s:步距
        '''
        inverted_residuals_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, 320, 1, 1]
        ]

        features = []
        #conv1 layer
        features.append(ConvBNReLU(3, input_channel, stride=2))

        for t, c, n, s in inverted_residuals_setting:
            output_channel = _make_divisible(c*alpha, round_nearest)
            for i in range(n):
                #每个n层bottleneck中只有第一层的stride=s, 剩余的层stride=1
                stride = s if i == 0 else 1
                features.append(block(input_channel, output_channel, stride, expand_ration=t))
                input_channel = output_channel

        features.append(ConvBNReLU(input_channel, last_channel, 1))

        self.features = nn.Sequential(*features)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))

        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(last_channel, num_classes)
        )

        # 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.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)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x
相关推荐
Microvision维视智造10 分钟前
解析大尺寸液晶屏视觉检测,装配错位如何避免?
人工智能·计算机视觉·视觉检测
lilye6626 分钟前
精益数据分析(11/126):辨别虚荣指标,挖掘数据真价值
大数据·人工智能·数据分析
微学AI27 分钟前
详细介绍:MCP(大模型上下文协议)的架构与组件,以及MCP的开发实践
前端·人工智能·深度学习·架构·llm·mcp
Samuel-Gyx35 分钟前
2025第十六届蓝桥杯python B组满分题解(详细)
python·职场和发展·蓝桥杯
豆包MarsCode1 小时前
玩转MCP | 一文看懂如何在 Trae IDE 中解锁 MCP
人工智能·mcp·trae
行者无疆xcc1 小时前
【Django】设置让局域网内的人访问
后端·python·django
患得患失9491 小时前
【后端】【python】Python 爬虫常用的框架解析
开发语言·爬虫·python
我不是小upper1 小时前
详解机器学习各算法的优缺点!!
人工智能·算法·机器学习
小君1 小时前
New 版本Trea 对比 Cursor 选择你的下一代 AI 编程伙伴
前端·人工智能·trae
研一计算机小白一枚1 小时前
第一章:自然语言处理
人工智能·自然语言处理