def build_psr_fpn_backbone(cfg, input_shape: ShapeSpec, num_classes=None):
"""
Args:
cfg: a detectron2 CfgNode
Returns:
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
"""
bottom_up = build_resnet_backbone(cfg, input_shape)
in_features = cfg.MODEL.FPN.IN_FEATURES
out_channels = cfg.MODEL.FPN.OUT_CHANNELS
convf_name = cfg.MODEL.CUSTOM.FPN.CONVF_NAME
noise_var = cfg.MODEL.CUSTOM.FPN.NOISE_VAR
num_branch = cfg.MODEL.CUSTOM.BRANCH.NUM_BRANCH
backbone = PSRFPN(bottom_up=bottom_up,
in_features=in_features,
out_channels=out_channels,
convf_name=convf_name,
num_branch=num_branch,
noise_var=noise_var,
norm=cfg.MODEL.FPN.NORM,
top_block=LastLevelMaxPool(),
fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
num_classes=num_classes)
return backbone
def __init__(self, config: Config, *args, **kwargs):
super().__init__()
self.config = config
pretrained = config.get("pretrained", False)
pretrained_path = config.get("pretrained_path", None)
self.resnet = build_resnet_backbone(config, ShapeSpec(channels=3))
if pretrained:
state_dict = torch.hub.load_state_dict_from_url(pretrained_path,
progress=False)
new_state_dict = OrderedDict()
replace_layer = {"backbone.": ""}
for key, value in state_dict["model"].items():
new_key = re.sub(r"(backbone\.)",
lambda x: replace_layer[x.groups()[0]], key)
new_state_dict[new_key] = value
self.resnet.load_state_dict(new_state_dict, strict=False)
self.out_dim = 2048
def build_normal_fpn_pretrain_backbone(cfg,
input_shape: ShapeSpec,
num_classes=None):
"""
Args:
cfg: a detectron2 CfgNode
Returns:
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
"""
bottom_up = build_resnet_backbone(cfg, input_shape)
in_features = cfg.MODEL.FPN.IN_FEATURES
out_channels = cfg.MODEL.FPN.OUT_CHANNELS
noise_var = cfg.MODEL.CUSTOM.FPN.NOISE_VAR
backbone = FPN(bottom_up=bottom_up,
in_features=in_features,
out_channels=out_channels,
noise_var=noise_var,
norm=cfg.MODEL.FPN.NORM,
top_block=LastLevelMaxPool(),
fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
num_classes=num_classes)
return backbone