def modified_xception(pre_trained=True):
xception = AlignedXception(output_stride=16)
if pre_trained:
load_pretrained_xception("../pretrained_models/xception.pth", xception)
else:
xception = weight_initialization(xception)
return xception
def __init__(self, num_classes, initial='kaiming', backbone='resnet'):
super(DeebLab, self).__init__()
support_backbone = ['resnet', 'mobilenet', 'drn', 'vgg', 'xception']
assert support_backbone.__contains__(backbone), 'unsupported backbone'
if backbone == 'resnet':
self.backbone = resnet.resnet_101(pre_trained=True)
elif backbone == 'mobilenet':
self.backbone = mobilenet.mobilenet_v2(pre_trained=True)
elif backbone == 'drn':
self.backbone = drn.drn_d_105(pre_trained=True)
elif backbone == 'vgg':
self.backbone = vgg.vgg_16(pre_trained=True)
elif backbone == 'xception':
self.backbone = xception.modified_xception(pre_trained=True)
for m in self.backbone.modules():
if isinstance(m, nn.Conv2d):
self.in_dim = m.weight.data.size()[0]
self.dilation_rates = dilation_rates
self.padding_series = padding_series
self.classifier = weight_initialization(
ASPP_V2(in_dim=self.in_dim, dilation_rates=dilation_rates, padding_rates=padding_series,
num_classes=num_classes), init_type=initial)
# freeze bn layers
self.freeze_bn()
def vgg_16(pre_trained=True):
vgg = VGG()
if pre_trained:
load_pretrained_vgg("../pretrained_models/vgg_16.pth", vgg)
else:
vgg = weight_initialization(vgg)
return vgg
def resnet_101(pre_trained=True):
resnet = ResNet(BottleNeck, [3, 4, 23, 3])
if pre_trained:
load_pretrained_resnet(model_path, resnet)
else:
resnet = weight_initialization(resnet)
return resnet
def mobilenet_v2(pre_trained=True):
mbnet = MobileNetV2(block=InvertedResidual, output_stride=16)
if pre_trained:
load_pretrained_mobilenet(model_path, mbnet)
else:
mbnet = weight_initialization(mbnet)
return mbnet
def drn_d_105(pre_trained=True):
"""
drn_d_105
"""
drn = DRN(BottleNeck, [1, 1, 3, 4, 23, 3, 1, 1])
if pre_trained:
load_pretrained_drn("../pretrained_models/drn_d_105-12b40979.pth", drn)
else:
drn = weight_initialization(drn)
return drn
def shufflenet_v2_1x(pre_trained=True):
shufflenet = ShuffleNetV2(stage_repeat_nums=[4, 8, 4],
stage_out_dims=[24, 116, 232, 464, 1024],
num_classes=1000)
if pre_trained:
load_pretrained_shufflenet("../pretrained_models/shufflenetv2_x1.pth",
shufflenet)
else:
shufflenet = weight_initialization(shufflenet)
return shufflenet
def drn_d_105(pre_trained=True):
"""
drn_d_105
"""
drn = DRN(BottleNeck, [1, 1, 3, 4, 23, 3, 1, 1])
if pre_trained:
pretrained_dict = torch.load('../pretrained_models/drn_d_105.pth')
drn.load_state_dict(pretrained_dict)
else:
drn = weight_initialization(drn)
return drn