pytorch手动实现MobileNet_v2

First
在做项目的时候在GitHub上面找了一篇MobileV2模型实现的源码，自己仔细看了一下，感觉实现的只是整体结构，但是和论文种不太贴切，由此修改成较为符合论文结构的代码版本。

About原始代码
下面是原始代码实现BottleNeck的部分，可以看到只是实现了 升维->分组卷积->降维，并判断是否进行shortcut操作。
但是并没有具体到实现Depthwise Conv 和 Pointwise Conv两部分。

class LinearBottleNeck(nn.Module):

    def __init__(self, in_channels, out_channels, stride, t=6, class_num=100):
        super().__init__()

        self.residual = nn.Sequential(
            nn.Conv2d(in_channels, in_channels * t, 1),
            nn.BatchNorm2d(in_channels * t),
            nn.ReLU6(inplace=True),

            nn.Conv2d(in_channels * t, in_channels * t, 3, stride=stride, padding=1, groups=in_channels * t),
            nn.BatchNorm2d(in_channels * t),
            nn.ReLU6(inplace=True),

            nn.Conv2d(in_channels * t, out_channels, 1),
            nn.BatchNorm2d(out_channels)
        )

        self.stride = stride
        self.in_channels = in_channels
        self.out_channels = out_channels
    
    def forward(self, x):

        residual = self.residual(x)

        if self.stride == 1 and self.in_channels == self.out_channels:
            residual += x
        return residual

我打印了这个原始代码的模型输出的shape：
第一行是input shape，最后一行是output shape。

torch.Size([1, 3, 224, 224])
torch.Size([1, 32, 226, 226])
torch.Size([1, 16, 226, 226])
torch.Size([1, 24, 113, 113])
torch.Size([1, 32, 57, 57])
torch.Size([1, 64, 29, 29])
torch.Size([1, 96, 29, 29])
torch.Size([1, 160, 29, 29])
torch.Size([1, 320, 29, 29])
torch.Size([1, 100, 29, 29])
torch.Size([1, 100, 1, 1])
torch.Size([1, 100, 1, 1])
torch.Size([1, 100])

Second
About原理
参考：https://yinguobing.com/bottlenecks-block-in-mobilenetv2/
1Depthwise conv 和 Pointwise conv

Depthwise conv ：使用与输入通道数量相同数量的卷积核对input进行卷积，这里的卷积是一一对应的（一个卷积核对应一个channel 输出一个 Maps）
Pointwise conv ：使用大小维11input_channel的卷积核 对Depthwise conv 输出的结果进行卷积
一个卷积核卷积输出为一个Maps。

2Residual链接

当BottleNeck stride=1时 模型进行远跳连接，这种方式不改变模型的shape，即保持w，h，c 不变只进行特征变换。（这里的不改变是指featuremap input的shape 和 add 之后的 shape 相比）
当BottleNeck stride=2时，在Dwise会使featuremap的size缩小为1/2，channel也会根据设定进行改变。

3网络结构

t是在Dwise操作时 模型在提升维度时的参数。
c输出通道
n这个模块使用几次
s即stride
4网络结构
参考：https://blog.csdn.net/weixin_37918890/article/details/101298146

5Relu6函数的使用
在Pointwise conv时 和 在降维的时候 没有使用Relu6 激活函数 原因是会造成一些数据的损失。

Third
About 我的修改。
BottleNeck部分
首先使用1*1的卷积对featuremap进行升维，这里使用的是论文中的参数t ，把featuremap从input_channel升维到input_channel的维度

之后进行Dwise操作 （stride=2时 Depthwise conv使featuremap的size变为1/2）。
Pointwise conv不使用Relu6函数。

再使用1*1卷积把featuremap降维到BottleNeck 的ouput_channel 维度。这里的降维不使用Relu6激活函数，至少进行线性变换。

class LinearBottleNeck(nn.Module):
    def __init__(self, in_c, out_c, s, t):
        super().__init__()

        self.residual = nn.Sequential(
            #升维
            nn.Conv2d(in_c, in_c * t, 1),
            nn.BatchNorm2d(in_c * t),
            nn.ReLU6(inplace=True),

            # Depthwise
            nn.Conv2d(in_c * t, in_c * t, 3, stride=s, padding=1, groups=in_c * t),
            nn.BatchNorm2d(in_c * t),
            nn.ReLU6(inplace=True),

            # Pointwise
            nn.Conv2d(in_c * t, in_c * t, 1, stride=1, padding=0, groups=1),
            nn.BatchNorm2d(in_c * t),

            #降维
            nn.Conv2d(in_c * t, out_c, 1),
            nn.BatchNorm2d(out_c)
        )

        self.stride = s
        self.in_channels = in_c
        self.out_channels = out_c

    def forward(self, x):
        # print('before shortcut x shape is', x.shape)
        residual = self.residual(x)
        # ('before shortcut residual shape is', residual.shape)
        if self.stride == 1 and self.in_channels == self.out_channels:
            residual += x
        print('residual shape is', residual.shape)
        return residual

模型的每层输出的shape

---------input----------
torch.Size([1, 3, 224, 224])
---------pre----------
torch.Size([1, 32, 112, 112])
---------stage1----------
torch.Size([1, 16, 112, 112])
---------stage2----------
torch.Size([1, 24, 56, 56])
---------stage3----------
torch.Size([1, 32, 28, 28])
---------stage4----------
torch.Size([1, 64, 14, 14])
---------stage5----------
torch.Size([1, 96, 14, 14])
---------stage6----------
torch.Size([1, 160, 7, 7])
---------stage7----------
torch.Size([1, 320, 7, 7])
---------conv1----------
torch.Size([1, 1280, 7, 7])
---------adaptive_avg_pool----------
torch.Size([1, 1280, 1, 1])
---------conv2----------
torch.Size([1, 100, 1, 1])
---------ouput----------
torch.Size([1, 100])

模型结构打印（有些冗余）

MobileV2Net(
(pre): Sequential(
(0): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
)
(stage1): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(32, 32, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32)
(4): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(32, 32, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(stage2): Sequential(
(0): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(16, 96, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(96, 96, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=96)
(4): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(96, 96, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(96, 24, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(24, 144, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(144, 144, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=144)
(4): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(144, 144, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(144, 24, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
)
(stage3): Sequential(
(0): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(24, 144, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(144, 144, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=144)
(4): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(144, 144, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(144, 32, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=192)
(4): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(192, 192, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(192, 32, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(2): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=192)
(4): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(192, 192, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(192, 32, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
)
(stage4): Sequential(
(0): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(192, 192, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=192)
(4): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(192, 192, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(192, 64, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384)
(4): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(384, 384, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(2): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384)
(4): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(384, 384, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(3): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384)
(4): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(384, 384, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
)
(stage5): Sequential(
(0): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384)
(4): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(384, 384, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(384, 96, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(576, 576, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=576)
(4): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(576, 576, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(576, 96, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(2): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(576, 576, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=576)
(4): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(576, 576, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(576, 96, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
)
(stage6): Sequential(
(0): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(576, 576, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=576)
(4): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(576, 576, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(576, 160, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960)
(4): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(960, 960, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(2): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960)
(4): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(960, 960, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
)
(stage7): LinearBottleNeck(
(residual): Sequential(
(0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
(3): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960)
(4): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU6(inplace=True)
(6): Conv2d(960, 960, kernel_size=(1, 1), stride=(1, 1))
(7): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(8): Conv2d(960, 320, kernel_size=(1, 1), stride=(1, 1))
(9): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(conv1): Sequential(
(0): Conv2d(320, 1280, kernel_size=(1, 1), stride=(1, 1))
(1): BatchNorm2d(1280, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)
)
(conv2): Conv2d(1280, 100, kernel_size=(1, 1), stride=(1, 1))
)

代码

import torch
import torch.nn as nn
import numpy as np
from torchsummary import summary
# from padding_same_conv import Conv2d
import torch.nn.functional as F


class LinearBottleNeck(nn.Module):
    def __init__(self, in_c, out_c, s, t):
        super().__init__()

        self.residual = nn.Sequential(
            #升维
            nn.Conv2d(in_c, in_c * t, 1),
            nn.BatchNorm2d(in_c * t),
            nn.ReLU6(inplace=True),

            # Depthwise
            nn.Conv2d(in_c * t, in_c * t, 3, stride=s, padding=1, groups=in_c * t),
            nn.BatchNorm2d(in_c * t),
            nn.ReLU6(inplace=True),

            # Pointwise
            nn.Conv2d(in_c * t, in_c * t, 1, stride=1, padding=0, groups=1),
            nn.BatchNorm2d(in_c * t),

            #降维
            nn.Conv2d(in_c * t, out_c, 1),
            nn.BatchNorm2d(out_c)
        )

        self.stride = s
        self.in_channels = in_c
        self.out_channels = out_c

    def forward(self, x):
        # print('before shortcut x shape is', x.shape)
        residual = self.residual(x)
        # ('before shortcut residual shape is', residual.shape)
        if self.stride == 1 and self.in_channels == self.out_channels:
            residual += x
        # print('BottleNeck output shape is:', residual.shape)
        return residual


class MobileV2Net(nn.Module):
    def __init__(self, class_num=100):
        super().__init__()
        self.pre = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2,padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU6(inplace=True),
        )


        # Bottleneck repeat, inputs_channel, outputs_channel , s, t.
        self.stage1 = LinearBottleNeck(32, 16, 1, 1)
        self.stage2 = self.make_stage(2, 16, 24, 2, 6)
        self.stage3 = self.make_stage(3, 24, 32, 2, 6)
        self.stage4 = self.make_stage(4, 32, 64, 2, 6)
        self.stage5 = self.make_stage(3, 64, 96, 1, 6)
        self.stage6 = self.make_stage(3, 96, 160, 2, 6)
        self.stage7 = LinearBottleNeck(160, 320, 1, 6)

        self.conv1 = nn.Sequential(
            nn.Conv2d(320, 1280, 1),
            nn.BatchNorm2d(1280),
            nn.ReLU6(inplace=True)
        )

        self.conv2 = nn.Conv2d(1280, class_num, 1)


    def forward(self, x):
        print('---------input----------')
        print(x.shape)
        print('---------pre----------')
        x = self.pre(x)
        print(x.shape)
        print('---------stage1----------')
        x = self.stage1(x)
        print(x.shape)
        print('---------stage2----------')
        x = self.stage2(x)
        print(x.shape)
        print('---------stage3----------')
        x = self.stage3(x)
        print(x.shape)
        print('---------stage4----------')
        x = self.stage4(x)
        print(x.shape)
        print('---------stage5----------')
        x = self.stage5(x)
        print(x.shape)
        print('---------stage6----------')
        x = self.stage6(x)
        print(x.shape)
        print('---------stage7----------')
        x = self.stage7(x)
        print(x.shape)
        print('---------conv1----------')
        x = self.conv1(x)
        print(x.shape)
        print('---------adaptive_avg_pool----------')
        x = F.adaptive_avg_pool2d(x, 1)
        print(x.shape)
        print('---------conv2----------')
        x = self.conv2(x)
        print(x.shape)
        x = x.view(x.size(0), -1)
        print('---------ouput----------')
        print(x.shape)
        return x



    def make_stage(self, repeat, in_c, out_c, s, t):
        layers = []
        layers.append(LinearBottleNeck(in_c, out_c, s, t))

        while repeat - 1:
            layers.append(LinearBottleNeck(out_c, out_c, 1, t))
            repeat -= 1

        return nn.Sequential(*layers)

if __name__ == '__main__':
    model = MobileV2Net()
    print(model)

    a = torch.randn((1, 3, 224, 224))
    output = model(a)

本文地址：https://blog.csdn.net/weixin_41866216/article/details/107464637