-
Notifications
You must be signed in to change notification settings - Fork 2
/
mobilenetv2.py
131 lines (102 loc) · 4 KB
/
mobilenetv2.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
from functools import partial
import torch
import torch.nn as nn
from .ops import blocks
from .utils import export, config, load_from_local_or_url
from .ops.functional import make_divisible
from typing import Any, OrderedDict, List
@export
class MobileNetV2(nn.Module):
@blocks.activation(partial(nn.ReLU6, inplace=True))
def __init__(
self,
in_channels: int = 3,
num_classes: int = 1000,
multiplier: float = 1.0,
dropout_rate: float = 0.2,
dilations: List[int] = None,
thumbnail: bool = False,
**kwargs: Any
):
super().__init__()
dilations = [1] + (dilations or [1, 1, 1, 1])
assert len(dilations) == 5, ''
self.block = blocks.InvertedResidualBlock
FRONT_S = 1 if thumbnail else 2
t = [1, 6, 6, 6, 6, 6, 6]
c = [32, 16, 24, 32, 64, 96, 160, 320]
n = [1, 2, 3, 4, 3, 3, 1]
s = [1, FRONT_S, 2, 2, 1, 2, 1]
stages = [0, 1, 1, 1, 0, 1, 0]
if multiplier < 1.0:
c = [make_divisible(x * multiplier, 8) for x in c]
self.features = nn.Sequential(OrderedDict([
('stem', blocks.Stage(
blocks.Conv2dBlock(in_channels, c[0], 3, stride=FRONT_S)
))
]))
for i in range(len(t)):
layers = self.make_layers(
c[i],
t[i],
c[i+1],
n[i],
s[i],
dilations[len(self.features) + (stages[i] - 1)]
)
if stages[i]:
self.features.add_module(f'stage{len(self.features)}', blocks.Stage(layers))
else:
self.features[-1].append(layers)
self.features[-1].append(blocks.Conv2d1x1Block(c[-1], 1280))
self.pool = nn.AdaptiveAvgPool2d((1, 1))
self.classifier = nn.Sequential(
nn.Dropout(dropout_rate, inplace=True),
nn.Linear(1280, num_classes)
)
def make_layers(self, inp: int, t: int, oup: int, n: int, stride: int, dilation: int):
layers = [
self.block(
inp,
oup,
t,
stride=stride if dilation == 1 else 1,
dilation=max(dilation // stride, 1)
)
]
for _ in range(n - 1):
layers.append(self.block(oup, oup, t, dilation=dilation))
return layers
def forward(self, x):
x = self.features(x)
x = self.pool(x)
x = torch.flatten(x, 1)
x = self.classifier(x)
return x
def _mobilenet_v2(
multiplier: float = 1.0,
pretrained: bool = False,
pth: str = None,
progress: bool = True,
**kwargs: Any
):
model = MobileNetV2(multiplier=multiplier, **kwargs)
if pretrained:
load_from_local_or_url(model, pth, kwargs.get('url', None), progress)
return model
@export
@config(url='https://github.com/ffiirree/cv-models/releases/download/v0.0.1/mobilenet_v2_x1_0-bf342af4.pth')
def mobilenet_v2_x1_0(pretrained: bool = False, pth: str = None, progress: bool = True, **kwargs: Any):
return _mobilenet_v2(1.0, pretrained, pth, progress, **kwargs)
@export
@config(url='https://github.com/ffiirree/cv-models/releases/download/v0.0.1/mobilenet_v2_x0_75-fdfaf351.pth')
def mobilenet_v2_x0_75(pretrained: bool = False, pth: str = None, progress: bool = True, **kwargs: Any):
return _mobilenet_v2(0.75, pretrained, pth, progress, **kwargs)
@export
@config(url='https://github.com/ffiirree/cv-models/releases/download/v0.0.1/mobilenet_v2_x0_5-a9d4ed71.pth')
def mobilenet_v2_x0_5(pretrained: bool = False, pth: str = None, progress: bool = True, **kwargs: Any):
return _mobilenet_v2(0.5, pretrained, pth, progress, **kwargs)
@export
@config(url='https://github.com/ffiirree/cv-models/releases/download/v0.0.1/mobilenet_v2_x0_35-9bce1f31.pth')
def mobilenet_v2_x0_35(pretrained: bool = False, pth: str = None, progress: bool = True, **kwargs: Any):
return _mobilenet_v2(0.35, pretrained, pth, progress, **kwargs)