PyTorch——模型搭建——MobileNetV1（一）

MobileNetv1模型

PyTorch自带的MobileNetV1

没有实现MobileNetV1
详情参考：https://pytorch.org/vision/stable/models.html
自己搭建

 1 import torch
 2 import torch.nn as nn
 3 import torch.nn.functional as F
 4 
 5 class Block(nn.Module):
 6     "Depthwise conv + Pointwise conv"
 7     def __init__(self, in_channels, out_channels, stride=1):
 8         super(Block, self).__init__()
 9         self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=stride, padding=1, groups=in_channels, bias=False)
10         self.bn1 = nn.BatchNorm2d(in_channels)
11         self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False)
12         self.bn2 = nn.BatchNorm2d(out_channels)
13 
14     def forward(self, x):
15         x = self.conv1(x)
16         x = self.bn1(x)
17         x = F.relu(x)
18         x = self.conv2(x)
19         x = self.bn2(x)
20         x = F.relu(x)
21         return x
22 
23 class MobileNet(nn.Module):
24 
25     cfg = [64, (128, 2), 128, (256, 2), (512, 2), 512, 512, 512, 512, 512, (1024, 2), 1024]
26 
27     def __init__(self, num_classes=10):
28         super(MobileNet, self).__init__()
29         self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)
30         self.bn1 = nn.BatchNorm2d(32)
31         self.layers = self._make_layers(in_planes=32)
32         self.linear = nn.Linear(1024, num_classes)
33 
34     def _make_layers(self, in_planes):
35         layers = []
36         for x in self.cfg:
37             out_planes = x if isinstance(x, int) else x[0]
38             stride = 1 if isinstance(x, int) else x[1]
39             layers.append(Block(in_planes, out_planes, stride))
40             in_planes = out_planes
41         return nn.Sequential(*layers)
42 
43     def forward(self, x):
44         out = F.relu(self.bn1(self.conv1(x)))
45         out = self.layers(out)
46         out = F.avg_pool2d(out, 7)
47         out = out.view(out.size(0), -1)
48         out = self.linear(out)
49         return out
50 
51 input = torch.randn(32, 3, 224, 224)
52 net = MobileNet(8)
53 out = net(input)
54 print(out.size())