mirror
/
MODNet
mirror of https://github.com/ZHKKKe/MODNet.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add torchscript

pull/78/head
kevin 2021-03-05 14:23:49 +08:00
parent e481792869
commit 57104c242e
2 changed files with 48 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
src/models/backbones/mobilenetv2.py
View File

@ -136,17 +136,31 @@ class MobileNetV2(nn.Module):
# Initialize weights
self._init_weights()
def forward(self, x, feature_names=None):
def forward(self, x):
# Stage1
x = reduce(lambda x, n: self.features[n](x), list(range(0,2)), x)
x = self.features[0](x)
x = self.features[1](x)
# Stage2
x = reduce(lambda x, n: self.features[n](x), list(range(2,4)), x)
x = self.features[2](x)
x = self.features[3](x)
# Stage3
x = reduce(lambda x, n: self.features[n](x), list(range(4,7)), x)
x = self.features[4](x)
x = self.features[5](x)
x = self.features[6](x)
# Stage4
x = reduce(lambda x, n: self.features[n](x), list(range(7,14)), x)
x = self.features[7](x)
x = self.features[8](x)
x = self.features[9](x)
x = self.features[10](x)
x = self.features[11](x)
x = self.features[12](x)
x = self.features[13](x)
# Stage5
x = reduce(lambda x, n: self.features[n](x), list(range(14,19)), x)
x = self.features[14](x)
x = self.features[15](x)
x = self.features[16](x)
x = self.features[17](x)
x = self.features[18](x)
# Classification
if self.num_classes is not None:

33
src/models/backbones/wrapper.py
View File

@ -36,15 +36,38 @@ class MobileNetV2Backbone(BaseBackbone):
self.enc_channels = [16, 24, 32, 96, 1280]
def forward(self, x):
x = reduce(lambda x, n: self.model.features[n](x), list(range(0, 2)), x)
# x = reduce(lambda x, n: self.model.features[n](x), list(range(0, 2)), x)
x = self.model.features[0](x)
x = self.model.features[1](x)
enc2x = x
x = reduce(lambda x, n: self.model.features[n](x), list(range(2, 4)), x)
# x = reduce(lambda x, n: self.model.features[n](x), list(range(2, 4)), x)
x = self.model.features[2](x)
x = self.model.features[3](x)
enc4x = x
x = reduce(lambda x, n: self.model.features[n](x), list(range(4, 7)), x)
# x = reduce(lambda x, n: self.model.features[n](x), list(range(4, 7)), x)
x = self.model.features[4](x)
x = self.model.features[5](x)
x = self.model.features[6](x)
enc8x = x
x = reduce(lambda x, n: self.model.features[n](x), list(range(7, 14)), x)
# x = reduce(lambda x, n: self.model.features[n](x), list(range(7, 14)), x)
x = self.model.features[7](x)
x = self.model.features[8](x)
x = self.model.features[9](x)
x = self.model.features[10](x)
x = self.model.features[11](x)
x = self.model.features[12](x)
x = self.model.features[13](x)
enc16x = x
x = reduce(lambda x, n: self.model.features[n](x), list(range(14, 19)), x)
# x = reduce(lambda x, n: self.model.features[n](x), list(range(14, 19)), x)
x = self.model.features[14](x)
x = self.model.features[15](x)
x = self.model.features[16](x)
x = self.model.features[17](x)
x = self.model.features[18](x)
enc32x = x
return [enc2x, enc4x, enc8x, enc16x, enc32x]