def __init__(self, layer):
# yapf: disable
layers = [
Parallel([nn.Identity(), layer]),
Sum()
]
# yapf: enable
super().__init__(*layers)
def __init__(self,
in_feats_shapes: Sequence[Tuple[int, ...]],
hidden_channels: int = 256,
out_channels: int = 2):
"""Constructor.
Args:
in_feats_shapes (Sequence[Tuple[int, ...]]): Shapes of the feature
maps that will be fed into the network. These are expected to
be tuples of the form (., C, H, ...).
hidden_channels (int, optional): The number of channels to which
all feature maps are convereted before being added together.
Defaults to 256.
out_channels (int, optional): Number of output channels. This will
normally be the number of classes. Defaults to 2.
"""
# reverse so that the deepest (i.e. produced by the deepest layer in
# the backbone network) feature map is first.
in_feats_shapes = in_feats_shapes[::-1]
in_feats_channels = [s[1] for s in in_feats_shapes]
# 1x1 conv to make the channels of all feature maps the same
in_convs = Parallel([
nn.Conv2d(in_channels, hidden_channels, kernel_size=1)
for in_channels in in_feats_channels
])
upsample_and_add = SequentialMultiInputMultiOutput(*[
Residual(
Interpolate(size=s[2:], mode='bilinear', align_corners=False))
for s in in_feats_shapes
])
out_convs = Parallel([
nn.Conv2d(hidden_channels, out_channels, kernel_size=3, padding=1)
for s in in_feats_shapes
])
# yapf: disable
layers = [
Reverse(),
in_convs,
upsample_and_add,
out_convs,
Reverse()
]
# yapf: enable
super().__init__(*layers)
def make_fused_backbone(old_backbone: nn.Module, new_backbone: nn.Module,
featureMapExtractorCls: Type,
channel_split: Tuple[int, int]) -> nn.Module:
"""Create a fused backbone using FuseNet style feature fusion.
See the paper, "FuseNet", by Hazirbas et al.,
https://vision.in.tum.de/_media/spezial/bib/hazirbasma2016fusenet.pdf.
"""
backbone = nn.Sequential(
SplitTensor(channel_split, dim=1),
Parallel([nn.Identity(),
featureMapExtractorCls(new_backbone)]),
featureMapExtractorCls(old_backbone, mode='fusion'))
return backbone
def _make_upsamplers(cls,
in_channels: int,
size: int,
num_upsamples_per_layer: Iterable[int],
num_groups: int = 32) -> Parallel:
layers = []
for num_upsamples in num_upsamples_per_layer:
upsampler = cls._upsample_feat(in_channels=in_channels,
num_upsamples=num_upsamples,
size=size,
num_groups=num_groups)
layers.append(upsampler)
upsamplers = Parallel(layers)
return upsamplers
def __init__(self, model: nn.Module, mode: Optional[str] = None):
super().__init__()
self.mode = mode
# yapf: disable
stem = nn.Sequential(
model.conv1,
model.bn1,
model.relu,
model.maxpool
)
layers = [
model.layer1,
model.layer2,
model.layer3,
model.layer4,
]
# yapf: enable
if mode == 'fusion':
self.m = nn.Sequential(
Parallel([stem, nn.Identity()]),
SequentialMultiInputMultiOutput(
*[nn.Sequential(Sum(), m) for m in layers]))
else:
self.m = SequentialMultiOutput(stem, *layers)
def __init__(self,
in_feats_shapes: Sequence[Tuple[int, ...]],
hidden_channels: int = 256,
out_channels: int = 2,
out_size: Optional[int] = None,
num_upsamples_per_layer: Optional[Sequence[int]] = None,
upsamplng_factor: int = 2,
num_groups_for_norm: int = 32):
"""Constructor.
Args:
in_feats_shapes (Sequence[Tuple[int, ...]]): Shapes of the feature
maps that will be fed into the network. These are expected to
be tuples of the form (., C, H, ...).
hidden_channels (int, optional): The number of channels to which
all feature maps are convereted before being added together.
Defaults to 256.
out_channels (int, optional): Number of output channels. This will
normally be the number of classes. Defaults to 2.
out_size (Optional[int], optional): Size of output. If None,
the size of the first feature map will be used.
Defaults to None.
num_upsamples_per_layer (Optional[Sequence[int]], optional): Number
of upsampling iterations for each feature map. Will depend on
the size of the feature map. Each upsampling iteration
comprises a conv-group_norm-relu block followed by a scaling
using torch.nn.functional.interpolate.
If None, each feature map is assumed to be half the size of the
preceeding one, meaning that it requires one more upsampling
iteration than the last one.
Defaults to None.
upsamplng_factor (int, optional): How much to scale per upsampling
iteration. Defaults to 2.
num_groups_for_norm (int, optional): Number of groups for group
norm layers. Defaults to 32.
"""
if num_upsamples_per_layer is None:
num_upsamples_per_layer = list(range(len(in_feats_shapes)))
if out_size is None:
out_size = in_feats_shapes[0][-2:]
in_convs = Parallel([
nn.Conv2d(s[1], hidden_channels, kernel_size=1)
for s in in_feats_shapes
])
upsamplers = self._make_upsamplers(
in_channels=hidden_channels,
size=out_size,
num_upsamples_per_layer=num_upsamples_per_layer,
num_groups=num_groups_for_norm)
out_conv = nn.Conv2d(hidden_channels // 2, out_channels, kernel_size=1)
# yapf: disable
layers = [
in_convs,
upsamplers,
Sum(),
out_conv
]
# yapf: enable
super().__init__(*layers)
def make_fpn_efficientnet(name: str = 'efficientnet_b0',
fpn_type: str = 'fpn',
out_size: Tuple[int, int] = (224, 224),
fpn_channels: int = 256,
num_classes: int = 1000,
pretrained: Optional[str] = 'imagenet',
in_channels: str = 3) -> nn.Module:
"""Loads the PyTorch implementation of EfficientNet from
https://github.com/lukemelas/EfficientNet-PyTorch using torch.hub.
Args:
name (str, optional): Name of the EfficientNet backbone. Only those
available in the lukemelas/EfficientNet-PyTorch repos are
supported. Defaults to 'efficientnet_b0'.
fpn_type (str, optional): Type of FPN. 'fpn' | 'panoptic' | 'panet'.
Defaults to 'fpn'.
out_size (Tuple[int, int], optional): Size of segmentation output.
Defaults to (224, 224).
fpn_channels (int, optional): Number of hidden channels to use in the
FPN. Defaults to 256.
num_classes (int, optional): Number of classes for which to make
predictions. Determines the channel width of the output.
Defaults to 1000.
pretrained (Optional[str], optional): One of
None | 'imagenet' | 'advprop'. See lukemelas/EfficientNet-PyTorch
for details. Defaults to True.
in_channels (int, optional): Channel width of the input. If greater
than 3, a parallel backbone is added to incorporate the new
channels and the feature maps of the two backbones are added
together to produce the final feature maps. Note that this is
currently different from make_fpn_resnet. See
lukemelas/EfficientNet-PyTorch for the in_channels < 3 case.
Defaults to 3.
Raises:
NotImplementedError: On unknown fpn_style.
Returns:
nn.Module: the FPN model
"""
effnet = _load_efficientnet(name=name,
num_classes=num_classes,
pretrained=pretrained)
if in_channels > 3:
new_channels = in_channels - 3
new_effnet = _load_efficientnet(
name=name,
num_classes=num_classes,
pretrained=pretrained,
in_channels=new_channels,
)
backbone = nn.Sequential(
SplitTensor((3, new_channels), dim=1),
Parallel([
EfficientNetFeatureMapsExtractor(effnet),
EfficientNetFeatureMapsExtractor(new_effnet)
]), AddAcross())
else:
backbone = EfficientNetFeatureMapsExtractor(effnet)
feat_shapes = _get_shapes(backbone, channels=in_channels, size=out_size)
if fpn_type == 'fpn':
fpn = nn.Sequential(
FPN(feat_shapes,
hidden_channels=fpn_channels,
out_channels=num_classes), SelectOne(idx=0))
elif fpn_type == 'panoptic':
fpn = PanopticFPN(feat_shapes,
hidden_channels=fpn_channels,
out_channels=num_classes)
elif fpn_type == 'panet+fpn':
feat_shapes2 = [(n, fpn_channels, h, w)
for (n, c, h, w) in feat_shapes]
fpn = nn.Sequential(
PANetFPN(feat_shapes,
hidden_channels=fpn_channels,
out_channels=fpn_channels),
FPN(feat_shapes2,
hidden_channels=fpn_channels,
out_channels=num_classes), SelectOne(idx=0))
else:
raise NotImplementedError()
# yapf: disable
model = nn.Sequential(
backbone,
fpn,
Interpolate(size=out_size, mode='bilinear', align_corners=False))
# yapf: enable
return model