backbone_utils.py 7.61 KB
Newer Older
1
import warnings
2

3
from torch import nn
4
from torchvision.ops import misc as misc_nn_ops
5
6
from torchvision.ops.feature_pyramid_network import FeaturePyramidNetwork, LastLevelMaxPool

7
from .. import mobilenet
8
from .. import resnet
9
from .._utils import IntermediateLayerGetter
10
11


eellison's avatar
eellison committed
12
class BackboneWithFPN(nn.Module):
13
14
15
16
    """
    Adds a FPN on top of a model.
    Internally, it uses torchvision.models._utils.IntermediateLayerGetter to
    extract a submodel that returns the feature maps specified in return_layers.
17
    The same limitations of IntermediateLayerGetter apply here.
18
    Args:
19
20
21
22
23
24
25
26
27
28
29
        backbone (nn.Module)
        return_layers (Dict[name, new_name]): a dict containing the names
            of the modules for which the activations will be returned as
            the key of the dict, and the value of the dict is the name
            of the returned activation (which the user can specify).
        in_channels_list (List[int]): number of channels for each feature map
            that is returned, in the order they are present in the OrderedDict
        out_channels (int): number of channels in the FPN.
    Attributes:
        out_channels (int): the number of channels in the FPN
    """
30

31
    def __init__(self, backbone, return_layers, in_channels_list, out_channels, extra_blocks=None):
eellison's avatar
eellison committed
32
        super(BackboneWithFPN, self).__init__()
33
34
35
36

        if extra_blocks is None:
            extra_blocks = LastLevelMaxPool()

eellison's avatar
eellison committed
37
38
        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
        self.fpn = FeaturePyramidNetwork(
39
40
            in_channels_list=in_channels_list,
            out_channels=out_channels,
41
            extra_blocks=extra_blocks,
42
43
44
        )
        self.out_channels = out_channels

eellison's avatar
eellison committed
45
46
47
48
49
    def forward(self, x):
        x = self.body(x)
        x = self.fpn(x)
        return x

50

51
52
53
54
55
56
def resnet_fpn_backbone(
    backbone_name,
    pretrained,
    norm_layer=misc_nn_ops.FrozenBatchNorm2d,
    trainable_layers=3,
    returned_layers=None,
57
    extra_blocks=None,
58
):
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
    """
    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]))]

78
    Args:
79
80
        backbone_name (string): resnet architecture. Possible values are 'ResNet', 'resnet18', 'resnet34', 'resnet50',
             'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d', 'wide_resnet50_2', 'wide_resnet101_2'
81
        pretrained (bool): If True, returns a model with backbone pre-trained on Imagenet
82
83
84
85
        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.
86
87
88
89
90
91
92
        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.
93
    """
94
    backbone = resnet.__dict__[backbone_name](pretrained=pretrained, norm_layer=norm_layer)
95

96
    # select layers that wont be frozen
97
    assert 0 <= trainable_layers <= 5
98
    layers_to_train = ["layer4", "layer3", "layer2", "layer1", "conv1"][:trainable_layers]
99
    if trainable_layers == 5:
100
        layers_to_train.append("bn1")
101
    for name, parameter in backbone.named_parameters():
102
        if all([not name.startswith(layer) for layer in layers_to_train]):
103
104
            parameter.requires_grad_(False)

105
106
107
108
109
110
    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
111
    return_layers = {f"layer{k}": str(v) for v, k in enumerate(returned_layers)}
112

113
    in_channels_stage2 = backbone.inplanes // 8
114
    in_channels_list = [in_channels_stage2 * 2 ** (i - 1) for i in returned_layers]
115
    out_channels = 256
116
    return BackboneWithFPN(backbone, return_layers, in_channels_list, out_channels, extra_blocks=extra_blocks)
117
118


119
def _validate_trainable_layers(pretrained, trainable_backbone_layers, max_value, default_value):
120
121
122
123
124
125
    # dont freeze any layers if pretrained model or backbone is not used
    if not pretrained:
        if trainable_backbone_layers is not None:
            warnings.warn(
                "Changing trainable_backbone_layers has not effect if "
                "neither pretrained nor pretrained_backbone have been set to True, "
126
127
                "falling back to trainable_backbone_layers={} so that all layers are trainable".format(max_value)
            )
128
129
130
        trainable_backbone_layers = max_value

    # by default freeze first blocks
131
    if trainable_backbone_layers is None:
132
133
        trainable_backbone_layers = default_value
    assert 0 <= trainable_backbone_layers <= max_value
134
    return trainable_backbone_layers
135
136
137
138
139
140
141
142
143


def mobilenet_backbone(
    backbone_name,
    pretrained,
    fpn,
    norm_layer=misc_nn_ops.FrozenBatchNorm2d,
    trainable_layers=2,
    returned_layers=None,
144
    extra_blocks=None,
145
146
147
):
    backbone = mobilenet.__dict__[backbone_name](pretrained=pretrained, norm_layer=norm_layer).features

148
    # Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks.
149
    # The first and last blocks are always included because they are the C0 (conv1) and Cn.
150
    stage_indices = [0] + [i for i, b in enumerate(backbone) if getattr(b, "_is_cn", False)] + [len(backbone) - 1]
151
152
153
154
    num_stages = len(stage_indices)

    # find the index of the layer from which we wont freeze
    assert 0 <= trainable_layers <= num_stages
155
    freeze_before = len(backbone) if trainable_layers == 0 else stage_indices[num_stages - trainable_layers]
156
157
158
159
160
161
162
163
164
165
166
167
168

    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
169
        return_layers = {f"{stage_indices[k]}": str(v) for v, k in enumerate(returned_layers)}
170
171
172
173
174
175
176
177
178
179
180

        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