def get_trunk(trunk_name, output_stride=8):
"""
Retrieve the network trunk and channel counts.
"""
assert output_stride == 8, 'Only stride8 supported right now'
if trunk_name == 'wrn38':
#
# FIXME: pass in output_stride once we support stride 16
#
backbone = wrn38(pretrained=True)
s2_ch = 128
s4_ch = 256
high_level_ch = 4096
elif trunk_name == 'seresnext-50' or trunk_name == 'seresnext-101':
backbone = get_resnet(trunk_name, output_stride=output_stride)
s2_ch = 48
s4_ch = -1
high_level_ch = 2048
elif trunk_name == 'resnet-50' or trunk_name == 'resnet-101':
backbone = get_resnet(trunk_name, output_stride=output_stride)
s2_ch = 256
s4_ch = -1
high_level_ch = 2048
elif trunk_name == 'hrnetv2':
backbone = hrnetv2.get_seg_model()
high_level_ch = backbone.high_level_ch
s2_ch = -1
s4_ch = -1
else:
raise 'unknown backbone {}'.format(trunk_name)
# logx.msg("Trunk: {}".format(trunk_name))
return backbone, s2_ch, s4_ch, high_level_ch
def get_trunk(trunk_name, output_stride=8):
"""
Retrieve the network trunk and channel counts.
"""
assert output_stride == 8, 'Only stride8 supported right now'
if trunk_name == 'wrn38':
#
# FIXME: pass in output_stride once we support stride 16
#
backbone = wrn38(pretrained=True)
s2_ch = 128
s4_ch = 256
high_level_ch = 4096
elif trunk_name == 'xception71':
backbone = xception71(output_stride=output_stride,
BatchNorm=Norm2d,
pretrained=True)
s2_ch = 64
s4_ch = 128
high_level_ch = 2048
elif trunk_name == 'seresnext-50' or trunk_name == 'seresnext-101':
backbone = get_resnet(trunk_name, output_stride=output_stride)
s2_ch = 48
s4_ch = -1
high_level_ch = 2048
elif trunk_name == 'resnet-50' or trunk_name == 'resnet-101':
backbone = get_resnet(trunk_name, output_stride=output_stride)
s2_ch = 256
s4_ch = -1
high_level_ch = 2048
elif trunk_name == 'hrnetv2':
backbone = hrnetv2.get_seg_model()
high_level_ch = backbone.high_level_ch
s2_ch = -1
s4_ch = -1
elif trunk_name == 'mobilenetv3_large':
backbone = MobileNetV3_Large(pretrained=False)
s2_ch = 16
s4_ch = 24
high_level_ch = 320
elif trunk_name == 'mobilenetv3_small':
backbone = MobileNetV3_Small(pretrained=False)
s2_ch = 16
s4_ch = 16
high_level_ch = 256
else:
raise 'unknown backbone {}'.format(trunk_name)
logx.msg("Trunk: {}".format(trunk_name))
return backbone, s2_ch, s4_ch, high_level_ch