MobileNet学习（二）——Pytorch搭建v1网络

先放上MobileNet v1的网络结构：


下面开始搭网络结构，主要分成三个部分，第一层的全连接、中间的深度可分离卷积以及最后的输出层，深度可分离卷积的层搭完后有一个平均池化的过程。最后，代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as tf

class MobileNetv1(nn.Module):
    parameters = [(32,64,1),(64,128,2),(128,128,1),(128,256,2),(256,256,1),
    (256,512,2),(512,512,1),(512,512,1),(512,512,1),(512,512,1),
    (512,512,1),(512,1024,2),(1024,1024,1)]

    def __init__(self,num_classes=1000):
        super(MobileNetv1,self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(
                in_channels=3,
                out_channels=32,
                kernel_size=3,
                stride=1,
                padding=1
            ),
            nn.BatchNorm2d(32),
            nn.ReLU(True),
        )
        self.conv_dw = self.Depthwise_Separable()
        self.linear = nn.Linear(1024,num_classes)

    def Depthwise_Separable(self):
        layers = []
        for parameter in self.parameters:
            in_channels = parameter[0]
            out_channels = parameter[1]
            stride = parameter[2]
            layer = nn.Sequential(
                nn.Conv2d(
                    in_channels=in_channels,
                    out_channels=in_channels,
                    kernel_size=3,
                    stride=stride,
                    padding=1,
                ),
                nn.BatchNorm2d(in_channels),
                nn.ReLU(True),
                nn.Conv2d(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    kernel_size=1,
                    stride=1,
                    padding=0,
                ),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(True),
            )
            layers.append(layer)
        return nn.Sequential(*layers,nn.AvgPool2d(7))

    def forward(self, x):
        x = self.conv(x)
        x = self.conv_dw(x)
        x = x.view(x.size(0), -1)
        x = self.linear(x)
        return x
 
net = MobileNetv1()
print(net)

打印出的网络结构如下：