def ResnestBackboneWithFPN():
backbone = timm.models.resnest.resnest50d(pretrained=True)
trainable_layers = 3
layers_to_train = ["layer4", "layer3", "layer2", "layer1",
"conv1"][:trainable_layers]
extra_blocks = LastLevelMaxPool()
# frextra_blocks eze layers only if pretrained backbone is used
for name, parameter in backbone.named_parameters():
if all([not name.startswith(layer) for layer in layers_to_train]):
parameter.requires_grad_(False)
returned_layers = [1, 2, 3, 4]
assert min(returned_layers) > 0 and max(returned_layers) < 5
return_layers = {
f"layer{k}": str(v)
for v, k in enumerate(returned_layers)
}
in_channels_list = [256, 512, 1024, 2048]
out_channels = 1024
return BackboneWithFPN(
backbone,
return_layers,
in_channels_list,
out_channels,
extra_blocks=LastLevelMaxPool(),
)
def resnet_fpn_backbone(
backbone_name,
pretrained,
norm_layer=misc_nn_ops.FrozenBatchNorm2d,
trainable_layers=3,
returned_layers=None,
extra_blocks=None
):
"""
Constructs a specified ResNet backbone with FPN on top. Freezes the specified number of layers in the backbone.
Examples::
>>> from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
>>> backbone = resnet_fpn_backbone('resnet50', pretrained=True, trainable_layers=3)
>>> # get some dummy image
>>> x = torch.rand(1,3,64,64)
>>> # compute the output
>>> output = backbone(x)
>>> print([(k, v.shape) for k, v in output.items()])
>>> # returns
>>> [('0', torch.Size([1, 256, 16, 16])),
>>> ('1', torch.Size([1, 256, 8, 8])),
>>> ('2', torch.Size([1, 256, 4, 4])),
>>> ('3', torch.Size([1, 256, 2, 2])),
>>> ('pool', torch.Size([1, 256, 1, 1]))]
Args:
backbone_name (string): resnet architecture. Possible values are 'ResNet', 'resnet18', 'resnet34', 'resnet50',
'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d', 'wide_resnet50_2', 'wide_resnet101_2'
norm_layer (torchvision.ops): it is recommended to use the default value. For details visit:
(https://github.com/facebookresearch/maskrcnn-benchmark/issues/267)
pretrained (bool): If True, returns a model with backbone pre-trained on Imagenet
trainable_layers (int): number of trainable (not frozen) resnet layers starting from final block.
Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.
"""
backbone = resnet.__dict__[backbone_name](
pretrained=pretrained,
norm_layer=norm_layer)
# select layers that wont be frozen
assert 0 <= trainable_layers <= 5
layers_to_train = ['layer4', 'layer3', 'layer2', 'layer1', 'conv1'][:trainable_layers]
# freeze layers only if pretrained backbone is used
for name, parameter in backbone.named_parameters():
if all([not name.startswith(layer) for layer in layers_to_train]):
parameter.requires_grad_(False)
if extra_blocks is None:
extra_blocks = LastLevelMaxPool()
if returned_layers is None:
returned_layers = [1, 2, 3, 4]
assert min(returned_layers) > 0 and max(returned_layers) < 5
return_layers = {f'layer{k}': str(v) for v, k in enumerate(returned_layers)}
in_channels_stage2 = backbone.inplanes // 8
in_channels_list = [in_channels_stage2 * 2 ** (i - 1) for i in returned_layers]
out_channels = 256
return BackboneWithFPN(backbone, return_layers, in_channels_list, out_channels, extra_blocks=extra_blocks)
def __init__(self, config):
super(Faster_Rcnn, self).__init__()
self.config = config
self.Mean = torch.tensor(config.Mean, dtype=torch.float32)
self.num_anchor = len(config.anchor_scales) * len(config.anchor_ratios)
self.anchors = []
self.num_anchor = []
for i in range(5):
self.num_anchor.append(len(config.anchor_scales[i]) * len(config.anchor_ratios[i]))
stride = 4 * 2 ** i
print(stride, self.config.anchor_scales[i], self.config.anchor_ratios[i])
anchors = get_anchors(np.ceil(self.config.img_max / stride + 1), self.config.anchor_scales[i],
self.config.anchor_ratios[i], stride=stride)
print(anchors.shape)
self.anchors.append(anchors)
self.PC = ProposalCreator(nms_thresh=config.roi_nms_thresh,
n_train_pre_nms=config.roi_train_pre_nms,
n_train_post_nms=config.roi_train_post_nms,
n_test_pre_nms=config.roi_test_pre_nms, n_test_post_nms=config.roi_test_post_nms,
min_size=config.roi_min_size)
self.features = resnet50()
self.fpn = FPN_net([256, 512, 1024, 2048], 256, extra_blocks=LastLevelMaxPool())
self.rpn = RPN_net(256, self.num_anchor[0])
self.roialign = MultiScaleRoIAlign(['feat0', 'feat1', 'feat2', 'feat3'], 7, 2)
self.fast = Fast_net(config.num_cls, 256 * 7 * 7, 1024)
self.a = 0
self.b = 0
self.c = 0
self.d = 0
self.fast_num = 0
self.fast_num_P = 0
def _mobilenet_extractor(
backbone: Union[mobilenet.MobileNetV2, mobilenet.MobileNetV3],
fpn: bool,
trainable_layers: int,
returned_layers: Optional[List[int]] = None,
extra_blocks: Optional[ExtraFPNBlock] = None,
) -> nn.Module:
backbone = backbone.features
# Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks.
# The first and last blocks are always included because they are the C0 (conv1) and Cn.
stage_indices = [0] + [
i for i, b in enumerate(backbone) if getattr(b, "_is_cn", False)
] + [len(backbone) - 1]
num_stages = len(stage_indices)
# find the index of the layer from which we wont freeze
if trainable_layers < 0 or trainable_layers > num_stages:
raise ValueError(
f"Trainable layers should be in the range [0,{num_stages}], got {trainable_layers} "
)
freeze_before = len(backbone) if trainable_layers == 0 else stage_indices[
num_stages - trainable_layers]
for b in backbone[:freeze_before]:
for parameter in b.parameters():
parameter.requires_grad_(False)
out_channels = 256
if fpn:
if extra_blocks is None:
extra_blocks = LastLevelMaxPool()
if returned_layers is None:
returned_layers = [num_stages - 2, num_stages - 1]
if min(returned_layers) < 0 or max(returned_layers) >= num_stages:
raise ValueError(
f"Each returned layer should be in the range [0,{num_stages - 1}], got {returned_layers} "
)
return_layers = {
f"{stage_indices[k]}": str(v)
for v, k in enumerate(returned_layers)
}
in_channels_list = [
backbone[stage_indices[i]].out_channels for i in returned_layers
]
return BackboneWithFPN(backbone,
return_layers,
in_channels_list,
out_channels,
extra_blocks=extra_blocks)
else:
m = nn.Sequential(
backbone,
# depthwise linear combination of channels to reduce their size
nn.Conv2d(backbone[-1].out_channels, out_channels, 1),
)
m.out_channels = out_channels # type: ignore[assignment]
return m
def __init__(self, backbone, return_layers, in_channels_list, out_channels):
super(BackboneWithFPN, self).__init__()
self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
self.fpn = FeaturePyramidNetwork(
in_channels_list=in_channels_list,
out_channels=out_channels,
extra_blocks=LastLevelMaxPool(),
)
self.out_channels = out_channels
def mobilenet_backbone(backbone_name,
pretrained,
fpn,
norm_layer=misc_nn_ops.FrozenBatchNorm2d,
trainable_layers=2,
returned_layers=None,
extra_blocks=None):
backbone = mobilenet.__dict__[backbone_name](
pretrained=pretrained, norm_layer=norm_layer).features
# Gather the indeces of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks.
# The first and last blocks are always included because they are the C0 (conv1) and Cn.
stage_indices = [0] + [
i for i, b in enumerate(backbone) if getattr(b, "is_strided", False)
] + [len(backbone) - 1]
num_stages = len(stage_indices)
# find the index of the layer from which we wont freeze
assert 0 <= trainable_layers <= num_stages
freeze_before = num_stages if trainable_layers == 0 else stage_indices[
num_stages - trainable_layers]
# freeze layers only if pretrained backbone is used
for b in backbone[:freeze_before]:
for parameter in b.parameters():
parameter.requires_grad_(False)
out_channels = 256
if fpn:
if extra_blocks is None:
extra_blocks = LastLevelMaxPool()
if returned_layers is None:
returned_layers = [num_stages - 2, num_stages - 1]
assert min(returned_layers) >= 0 and max(returned_layers) < num_stages
return_layers = {
f'{stage_indices[k]}': str(v)
for v, k in enumerate(returned_layers)
}
in_channels_list = [
backbone[stage_indices[i]].out_channels for i in returned_layers
]
return BackboneWithFPN(backbone,
return_layers,
in_channels_list,
out_channels,
extra_blocks=extra_blocks)
else:
m = nn.Sequential(
backbone,
# depthwise linear combination of channels to reduce their size
nn.Conv2d(backbone[-1].out_channels, out_channels, 1),
)
m.out_channels = out_channels
return m
def __init__(self, backbone):
super(TimmToVisionFPN, self).__init__()
self.backbone = backbone
self.out_channels = 256
self.in_channels_list = [128, 256, 512, 1024]
self.fpn = FeaturePyramidNetwork(
in_channels_list=self.in_channels_list,
out_channels=self.out_channels,
extra_blocks=LastLevelMaxPool(),
)
def __init__(self, backbone, return_layers, in_channels_list, out_channels):
body = IntermediateLayerGetter(backbone, return_layers=return_layers)
fpn = FeaturePyramidNetwork(
in_channels_list=in_channels_list,
out_channels=out_channels,
extra_blocks=LastLevelMaxPool(),
)
super(BackboneWithFPN, self).__init__(OrderedDict(
[("body", body), ("fpn", fpn)]))
self.out_channels = out_channels
def __init__(self, config):
super(Mask_Rcnn, self).__init__()
self.config = config
self.Mean = torch.tensor(config.Mean, dtype=torch.float32)
self.num_anchor = len(config.anchor_scales) * len(config.anchor_ratios)
self.anchors = []
self.num_anchor = []
for i in range(5):
self.num_anchor.append(
len(config.anchor_scales[i]) * len(config.anchor_ratios[i]))
stride = 4 * 2**i
print(stride, self.config.anchor_scales[i],
self.config.anchor_ratios[i])
anchors = get_anchors(np.ceil(self.config.img_max / stride + 1),
self.config.anchor_scales[i],
self.config.anchor_ratios[i],
stride=stride)
print(anchors.shape)
self.anchors.append(anchors)
self.ATC = AnchorTargetCreator(
n_sample=config.rpn_n_sample,
pos_iou_thresh=config.rpn_pos_iou_thresh,
neg_iou_thresh=config.rpn_neg_iou_thresh,
pos_ratio=config.rpn_pos_ratio)
self.PC = ProposalCreator(nms_thresh=config.roi_nms_thresh,
n_train_pre_nms=config.roi_train_pre_nms,
n_train_post_nms=config.roi_train_post_nms,
n_test_pre_nms=config.roi_test_pre_nms,
n_test_post_nms=config.roi_test_post_nms,
min_size=config.roi_min_size)
self.PTC = ProposalTargetCreator(
n_sample=config.fast_n_sample,
pos_ratio=config.fast_pos_ratio,
pos_iou_thresh=config.fast_pos_iou_thresh,
neg_iou_thresh_hi=config.fast_neg_iou_thresh_hi,
neg_iou_thresh_lo=config.fast_neg_iou_thresh_lo)
self.features = resnet101()
self.fpn = FPN_net([256, 512, 1024, 2048],
256,
extra_blocks=LastLevelMaxPool())
self.rpn = RPN_net(256, self.num_anchor[0])
self.roialign_7 = MultiScaleRoIAlign(
['feat0', 'feat1', 'feat2', 'feat3'], 7, 2)
self.roialign_14 = MultiScaleRoIAlign(
['feat0', 'feat1', 'feat2', 'feat3'], 14, 2)
self.roialign_28 = RoIAlign((28, 28), 1.0, 2)
self.fast = Fast_net(config.num_cls, 256 * 7 * 7, 1024)
self.mask_net = Mask_net(256, config.num_cls)
self.a = 0
self.b = 0
self.c = 0
self.d = 0
self.fast_num = 0
self.fast_num_P = 0
def __init__(self, backbone):
super(backboneWithFPN, self).__init__()
# tf_efficientnet_b4
self.body = IntermediateLayerGetter(backbone, return_layers={'block4': '1', 'block5': '2',
'block6': '3', 'conv_head': '4'})
self.fpn = FeaturePyramidNetwork(
in_channels_list=[160, 272, 448, 1792],
out_channels=256,
extra_blocks=LastLevelMaxPool(),
)
self.out_channels=256
def __init__(self, backbone):
super(TimmToVisionFPN, self).__init__()
self.backbone = backbone
self.out_channels = 256
## if you set timm model = resnet rather than cspresnet, you should set
## in_channels_list = [256, 512, 1024, 2048]
self.in_channels_list = [128, 256, 512, 1024]
self.fpn = FeaturePyramidNetwork(
in_channels_list=self.in_channels_list,
out_channels=self.out_channels,
extra_blocks=LastLevelMaxPool(),
)
def resnet_fpn_backbone(
backbone_name,
pretrained,
norm_layer=misc_nn_ops.FrozenBatchNorm2d,
trainable_layers=3,
returned_layers=None,
extra_blocks=None
):
"""
Constructs a specified ResNet backbone with FPN on top. Freezes the specified number of layers in the backbone.
Args:
backbone_name (string): resnet architecture. Possible values are 'ResNet', 'resnet18', 'resnet34', 'resnet50',
'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d', 'wide_resnet50_2', 'wide_resnet101_2'
pretrained (bool): If True, returns a model with backbone pre-trained on Imagenet
norm_layer (torchvision.ops): it is recommended to use the default value. For details visit:
(https://github.com/facebookresearch/maskrcnn-benchmark/issues/267)
trainable_layers (int): number of trainable (not frozen) resnet layers starting from final block.
Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.
returned_layers (list of int): The layers of the network to return. Each entry must be in ``[1, 4]``.
By default all layers are returned.
extra_blocks (ExtraFPNBlock or None): if provided, extra operations will
be performed. It is expected to take the fpn features, the original
features and the names of the original features as input, and returns
a new list of feature maps and their corresponding names. By
default a ``LastLevelMaxPool`` is used.
"""
backbone = resnet.__dict__[backbone_name](
pretrained=pretrained,
norm_layer=norm_layer)
# select layers that wont be frozen
assert 0 <= trainable_layers <= 5
layers_to_train = ['layer4', 'layer3', 'layer2', 'layer1', 'conv1'][:trainable_layers]
if trainable_layers == 5:
layers_to_train.append('bn1')
for name, parameter in backbone.named_parameters():
if all([not name.startswith(layer) for layer in layers_to_train]):
parameter.requires_grad_(False)
if extra_blocks is None:
extra_blocks = LastLevelMaxPool()
if returned_layers is None:
returned_layers = [1, 2, 3, 4]
assert min(returned_layers) > 0 and max(returned_layers) < 5
return_layers = {f'layer{k}': str(v) for v, k in enumerate(returned_layers)}
in_channels_stage2 = backbone.inplanes // 8
in_channels_list = [in_channels_stage2 * 2 ** (i - 1) for i in returned_layers]
out_channels = 256
return BackboneWithFPN(backbone, return_layers, in_channels_list, out_channels, extra_blocks=extra_blocks)
def __init__(self, backbone, return_layers, in_channels_list, out_channels,
ext_config):
super().__init__()
if ext_config.get('backbone_frozen', False):
module_util.freeze_module_params(backbone)
self.body = ExtIntermediateLayerGetter(backbone,
return_layers=return_layers,
ext_config=ext_config)
self.fpn = FeaturePyramidNetwork(
in_channels_list=in_channels_list,
out_channels=out_channels,
extra_blocks=LastLevelMaxPool(),
)
self.out_channels = out_channels
self.split = False
def mobilenet_backbone(
backbone_name: str,
pretrained: bool,
norm_layer: Callable[..., nn.Module] = misc_nn_ops.FrozenBatchNorm2d,
trainable_layers: int = 2,
returned_layers: Optional[List[int]] = None,
) -> nn.Module:
backbone = mobilenet.__dict__[backbone_name](
pretrained=pretrained, norm_layer=norm_layer).features
# Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks.
# The first and last blocks are always included because they are the C0 (conv1) and Cn.
stage_indices = (
[0] +
[i for i, b in enumerate(backbone) if getattr(b, "_is_cn", False)] +
[len(backbone) - 1])
num_stages = len(stage_indices)
# find the index of the layer from which we wont freeze
assert 0 <= trainable_layers <= num_stages
freeze_before = len(backbone) if trainable_layers == 0 else stage_indices[
num_stages - trainable_layers]
for b in backbone[:freeze_before]:
for parameter in b.parameters():
parameter.requires_grad_(False)
out_channels = 256
if returned_layers is None:
returned_layers = [num_stages - 3, num_stages - 2, num_stages - 1]
assert min(returned_layers) >= 0 and max(returned_layers) < num_stages
return_layers = {
f"{stage_indices[k]}": str(v)
for v, k in enumerate(returned_layers)
}
in_channels_list = [
backbone[stage_indices[i]].out_channels for i in returned_layers
]
return BackboneWithFPN(
backbone,
return_layers,
in_channels_list,
out_channels,
extra_blocks=LastLevelMaxPool(),
)
def maskrcnn_resnet18_fpn(num_classes):
src_backbone = torchvision.models.resnet18(pretrained=True)
# 去掉后面的全连接层
return_layers = {'layer1': 0,
'layer2': 1,
'layer3': 2,
'layer4': 3}
names = [name for name, _ in src_backbone.named_children()]
# just 验证，失败则报异常
if not set(return_layers).issubset(names):
raise ValueError("return_layers are not present in model")
orig_return_layers = return_layers
# 复制一份到 layers
return_layers = {k: v for k, v in return_layers.items()}
layers = OrderedDict()
for name, module in src_backbone.named_children():
layers[name] = module
if name in return_layers:
del return_layers[name]
if not return_layers:
break
# 得到去掉池化、全连接层的模型
backbone_module = backbone_body(layers, orig_return_layers)
# FPN层，resnet18 layer4 chanels为 512，fpn顶层512/8
in_channels_stage2 = 64
in_channels_list = [
in_channels_stage2,
in_channels_stage2 * 2,
in_channels_stage2 * 4,
in_channels_stage2 * 8,
]
out_channels = 64
fpn = FeaturePyramidNetwork(
in_channels_list=in_channels_list,
out_channels=out_channels,
extra_blocks=LastLevelMaxPool(),
)
backbone_fpn = BackboneFPN(backbone_module,
fpn,
out_channels)
model = MaskRCNN(backbone_fpn, num_classes)
return model
def __init__(self, config):
super(Mask_Rcnn, self).__init__()
self.config = config
self.Mean = torch.tensor(config.Mean, dtype=torch.float32)
self.num_anchor = len(config.anchor_scales) * len(config.anchor_ratios)
self.anchors = []
self.num_anchor = []
for i in range(5):
self.num_anchor.append(
len(config.anchor_scales[i]) * len(config.anchor_ratios[i]))
stride = 4 * 2**i
print(stride, self.config.anchor_scales[i],
self.config.anchor_ratios[i])
anchors = get_anchors(np.ceil(self.config.img_max / stride + 1),
self.config.anchor_scales[i],
self.config.anchor_ratios[i],
stride=stride)
print(anchors.shape)
self.anchors.append(anchors)
self.PC = ProposalCreator(nms_thresh=config.roi_nms_thresh,
n_train_pre_nms=config.roi_train_pre_nms,
n_train_post_nms=config.roi_train_post_nms,
n_test_pre_nms=config.roi_test_pre_nms,
n_test_post_nms=config.roi_test_post_nms,
min_size=config.roi_min_size)
self.features = resnet50()
self.fpn = FPN_net([256, 512, 1024, 2048],
256,
extra_blocks=LastLevelMaxPool())
self.rpn = RPN_net(256, self.num_anchor[0])
self.roialign_7 = MultiScaleRoIAlign(
['feat0', 'feat1', 'feat2', 'feat3'], 7, 2)
self.roialign_14 = MultiScaleRoIAlign(
['feat0', 'feat1', 'feat2', 'feat3'], 14, 2)
# self.roialign_28 = RoIAlign((28, 28), 1.0, 2)
self.fast = Fast_net(config.num_cls, 256 * 7 * 7, 1024)
self.fast_2 = Fast_net(config.num_cls, 256 * 7 * 7, 1024)
self.fast_3 = Fast_net(config.num_cls, 256 * 7 * 7, 1024)
self.mask_net = Mask_net(256, config.num_cls)
self.loc_std1 = [1. / 10, 1. / 10, 1. / 5, 1. / 5]
self.loc_std2 = [1. / 20, 1. / 20, 1. / 10, 1. / 10]
self.loc_std3 = [1. / 30, 1. / 30, 1. / 15, 1. / 15]
self.weights = [1.0, 1.0, 1.0]
def _resnet_fpn_extractor(
backbone: resnet.ResNet,
trainable_layers: int,
returned_layers: Optional[List[int]] = None,
extra_blocks: Optional[ExtraFPNBlock] = None,
norm_layer: Optional[Callable[..., nn.Module]] = None,
) -> BackboneWithFPN:
# select layers that wont be frozen
if trainable_layers < 0 or trainable_layers > 5:
raise ValueError(
f"Trainable layers should be in the range [0,5], got {trainable_layers}"
)
layers_to_train = ["layer4", "layer3", "layer2", "layer1",
"conv1"][:trainable_layers]
if trainable_layers == 5:
layers_to_train.append("bn1")
for name, parameter in backbone.named_parameters():
if all([not name.startswith(layer) for layer in layers_to_train]):
parameter.requires_grad_(False)
if extra_blocks is None:
extra_blocks = LastLevelMaxPool()
if returned_layers is None:
returned_layers = [1, 2, 3, 4]
if min(returned_layers) <= 0 or max(returned_layers) >= 5:
raise ValueError(
f"Each returned layer should be in the range [1,4]. Got {returned_layers}"
)
return_layers = {
f"layer{k}": str(v)
for v, k in enumerate(returned_layers)
}
in_channels_stage2 = backbone.inplanes // 8
in_channels_list = [
in_channels_stage2 * 2**(i - 1) for i in returned_layers
]
out_channels = 256
return BackboneWithFPN(backbone,
return_layers,
in_channels_list,
out_channels,
extra_blocks=extra_blocks,
norm_layer=norm_layer)
def lambda_resnet_fpn_backbone(backbone_name,
pretrained,
norm_layer=misc_nn_ops.FrozenBatchNorm2d,
trainable_layers=5,
returned_layers=None,
extra_blocks=None):
backbone = lambda_resnet.__dict__[backbone_name](
# pretrained=pretrained,
# norm_layer=norm_layer
)
# select layers that wont be frozen
assert trainable_layers <= 5 and trainable_layers >= 0
layers_to_train = ['layer4', 'layer3', 'layer2', 'layer1',
'conv1'][:trainable_layers]
# freeze layers only if pretrained backbone is used
for name, parameter in backbone.named_parameters():
if all([not name.startswith(layer) for layer in layers_to_train]):
parameter.requires_grad_(False)
if extra_blocks is None:
extra_blocks = LastLevelMaxPool()
if returned_layers is None:
returned_layers = [1, 2, 3, 4]
assert min(returned_layers) > 0 and max(returned_layers) < 5
return_layers = {
f'layer{k}': str(v)
for v, k in enumerate(returned_layers)
}
in_channels_stage2 = backbone.inplanes // 8
in_channels_list = [
in_channels_stage2 * 2**(i - 1) for i in returned_layers
]
out_channels = 256
return BackboneWithFPN(backbone,
return_layers,
in_channels_list,
out_channels,
extra_blocks=extra_blocks)
def efficientnet_fpn_backbone(backbone_name: str = "efficientnet-b0",
return_blocks: Optional[List[int]] = None,
out_channels: int = 256):
efficient_net: EfficientNet = EfficientNet.from_pretrained(backbone_name)
return_blocks = return_blocks if return_blocks is not None else BACKBONE_TO_RETURN_BLOCKS.get(
backbone_name, [])
backbone = EfficientNetBackbone(efficient_net, return_blocks)
fpn = FeaturePyramidNetwork(
in_channels_list=[
efficient_net._blocks[rb]._block_args.output_filters
for rb in return_blocks
],
out_channels=out_channels,
extra_blocks=LastLevelMaxPool(),
)
result = nn.Sequential(backbone, fpn)
result.out_channels = out_channels
return result
def __init__(self, backbone):
super(backboneWithFPN, self).__init__()
# resnet
# self.body = IntermediateLayerGetter(backbone, return_layers={'layer2': '1', 'layer3': '2', 'layer4': '3'})
# efficientnet
self.body = IntermediateLayerGetter(
backbone,
return_layers={
"block3": "0",
"block4": "1",
"block5": "2",
"block6": "3",
},
)
self.fpn = FeaturePyramidNetwork(
in_channels_list=[160, 224, 384, 640],
out_channels=160,
extra_blocks=LastLevelMaxPool(),
)
self.out_channels = 160
def __init__(
self,
backbone: nn.Module,
return_layers: Dict[str, str],
in_channels_list: List[int],
out_channels: int,
extra_blocks: Optional[ExtraFPNBlock] = None,
) -> None:
super().__init__()
if extra_blocks is None:
extra_blocks = LastLevelMaxPool()
self.body = IntermediateLayerGetter(backbone,
return_layers=return_layers)
self.fpn = FeaturePyramidNetwork(
in_channels_list=in_channels_list,
out_channels=out_channels,
extra_blocks=extra_blocks,
)
self.out_channels = out_channels
def _resnet_fpn_extractor(
backbone: resnet.ResNet,
trainable_layers: int,
returned_layers: Optional[List[int]] = None,
extra_blocks: Optional[ExtraFPNBlock] = None,
) -> BackboneWithFPN:
# select layers that wont be frozen
assert 0 <= trainable_layers <= 5
layers_to_train = ["layer4", "layer3", "layer2", "layer1",
"conv1"][:trainable_layers]
if trainable_layers == 5:
layers_to_train.append("bn1")
for name, parameter in backbone.named_parameters():
if all([not name.startswith(layer) for layer in layers_to_train]):
parameter.requires_grad_(False)
if extra_blocks is None:
extra_blocks = LastLevelMaxPool()
if returned_layers is None:
returned_layers = [1, 2, 3, 4]
assert min(returned_layers) > 0 and max(returned_layers) < 5
return_layers = {
f"layer{k}": str(v)
for v, k in enumerate(returned_layers)
}
in_channels_stage2 = backbone.inplanes // 8
in_channels_list = [
in_channels_stage2 * 2**(i - 1) for i in returned_layers
]
out_channels = 256
return BackboneWithFPN(backbone,
return_layers,
in_channels_list,
out_channels,
extra_blocks=extra_blocks)
def my_resnet_fpn_backbone(backbone,
pretrained,
norm_layer=misc_nn_ops.FrozenBatchNorm2d,
trainable_layers=3,
returned_layers=None,
extra_blocks=None):
# select layers that wont be frozen
assert 0 <= trainable_layers <= 5
layers_to_train = ['layer4', 'layer3', 'layer2', 'layer1',
'conv1'][:trainable_layers]
if trainable_layers == 5:
layers_to_train.append('bn1')
for name, parameter in backbone.named_parameters():
if all([not name.startswith(layer) for layer in layers_to_train]):
parameter.requires_grad_(False)
if extra_blocks is None:
extra_blocks = LastLevelMaxPool()
if returned_layers is None:
returned_layers = [1, 2, 3, 4]
assert min(returned_layers) > 0 and max(returned_layers) < 5
return_layers = {
f'layer{k}': str(v)
for v, k in enumerate(returned_layers)
}
in_channels_stage2 = backbone.inplanes // 8
in_channels_list = [
in_channels_stage2 * 2**(i - 1) for i in returned_layers
]
out_channels = 256
return BackboneWithFPN(backbone,
return_layers,
in_channels_list,
out_channels,
extra_blocks=extra_blocks)
def __init__(self,
backbone,
return_layers,
in_channels_list,
out_channels,
last_level='pool'):
body = IntermediateLayerGetter(backbone, return_layers=return_layers)
if last_level == 'pool':
extra_blocks = LastLevelMaxPool()
elif last_level == 'p6p7':
extra_blocks = LastLevelP6P7(in_channels_list[-1], out_channels)
else:
raise ValueError(f'Not implemented {last_level}.')
fpn = FeaturePyramidNetwork(
in_channels_list=in_channels_list,
out_channels=out_channels,
extra_blocks=extra_blocks,
)
super(BackboneWithFPN,
self).__init__(OrderedDict([("body", body), ("fpn", fpn)]))
self.out_channels = out_channels
def __init__(self, in_channels, out_channels=256, weight_file=None):
super(TorchFPN, self).__init__()
self.weight_file = weight_file
self.fpn = FeaturePyramidNetwork(in_channels_list=in_channels,
out_channels=out_channels,
extra_blocks=LastLevelMaxPool())
def __init__(self, in_channels, out_channels=256):
super(FPN, self).__init__(in_channels_list=in_channels,
out_channels=out_channels,
extra_blocks=LastLevelMaxPool())