File size: 2,610 Bytes
206ce41 |
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 |
import os
from functools import reduce
import torch
import torch.nn as nn
from .mobilenetv2 import MobileNetV2
class BaseBackbone(nn.Module):
""" Superclass of Replaceable Backbone Model for Semantic Estimation
"""
def __init__(self, in_channels):
super(BaseBackbone, self).__init__()
self.in_channels = in_channels
self.model = None
self.enc_channels = []
def forward(self, x):
raise NotImplementedError
def load_pretrained_ckpt(self):
raise NotImplementedError
class MobileNetV2Backbone(BaseBackbone):
""" MobileNetV2 Backbone
"""
def __init__(self, in_channels):
super(MobileNetV2Backbone, self).__init__(in_channels)
self.model = MobileNetV2(self.in_channels, alpha=1.0, expansion=6, num_classes=None)
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 = 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 = 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 = 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 = 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 = 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]
def load_pretrained_ckpt(self):
# the pre-trained model is provided by https://github.com/thuyngch/Human-Segmentation-PyTorch
ckpt_path = './pretrained/mobilenetv2_human_seg.ckpt'
if not os.path.exists(ckpt_path):
print('cannot find the pretrained mobilenetv2 backbone')
exit()
ckpt = torch.load(ckpt_path)
self.model.load_state_dict(ckpt)
|