resnext.py 7.41 KB
Newer Older
pangjm's avatar
pangjm committed
1
2
3
4
import math

import torch.nn as nn

Kai Chen's avatar
Kai Chen committed
5
from ..registry import BACKBONES
6
from ..utils import build_conv_layer, build_norm_layer
7
8
from .resnet import Bottleneck as _Bottleneck
from .resnet import ResNet
pangjm's avatar
pangjm committed
9
10


pangjm's avatar
pangjm committed
11
class Bottleneck(_Bottleneck):
pangjm's avatar
pangjm committed
12

Kai Chen's avatar
Kai Chen committed
13
    def __init__(self, inplanes, planes, groups=1, base_width=4, **kwargs):
pangjm's avatar
pangjm committed
14
        """Bottleneck block for ResNeXt.
pangjm's avatar
pangjm committed
15
16
17
        If style is "pytorch", the stride-two layer is the 3x3 conv layer,
        if it is "caffe", the stride-two layer is the first 1x1 conv layer.
        """
Kai Chen's avatar
Kai Chen committed
18
        super(Bottleneck, self).__init__(inplanes, planes, **kwargs)
pangjm's avatar
pangjm committed
19

pangjm's avatar
pangjm committed
20
        if groups == 1:
pangjm's avatar
pangjm committed
21
            width = self.planes
pangjm's avatar
pangjm committed
22
        else:
pangjm's avatar
pangjm committed
23
            width = math.floor(self.planes * (base_width / 64)) * groups
pangjm's avatar
pangjm committed
24

yhcao6's avatar
yhcao6 committed
25
        self.norm1_name, norm1 = build_norm_layer(
Kai Chen's avatar
Kai Chen committed
26
            self.norm_cfg, width, postfix=1)
yhcao6's avatar
yhcao6 committed
27
        self.norm2_name, norm2 = build_norm_layer(
Kai Chen's avatar
Kai Chen committed
28
            self.norm_cfg, width, postfix=2)
yhcao6's avatar
yhcao6 committed
29
        self.norm3_name, norm3 = build_norm_layer(
Kai Chen's avatar
Kai Chen committed
30
            self.norm_cfg, self.planes * self.expansion, postfix=3)
ThangVu's avatar
ThangVu committed
31

32
33
        self.conv1 = build_conv_layer(
            self.conv_cfg,
pangjm's avatar
pangjm committed
34
35
36
37
38
            self.inplanes,
            width,
            kernel_size=1,
            stride=self.conv1_stride,
            bias=False)
39
        self.add_module(self.norm1_name, norm1)
yhcao6's avatar
yhcao6 committed
40
41
42
        fallback_on_stride = False
        self.with_modulated_dcn = False
        if self.with_dcn:
43
            fallback_on_stride = self.dcn.pop('fallback_on_stride', False)
yhcao6's avatar
yhcao6 committed
44
        if not self.with_dcn or fallback_on_stride:
45
46
            self.conv2 = build_conv_layer(
                self.conv_cfg,
yhcao6's avatar
yhcao6 committed
47
48
49
50
51
52
53
54
55
                width,
                width,
                kernel_size=3,
                stride=self.conv2_stride,
                padding=self.dilation,
                dilation=self.dilation,
                groups=groups,
                bias=False)
        else:
56
            assert self.conv_cfg is None, 'conv_cfg must be None for DCN'
57
58
            self.conv2 = build_conv_layer(
                self.dcn,
yhcao6's avatar
yhcao6 committed
59
60
61
62
63
64
65
66
                width,
                width,
                kernel_size=3,
                stride=self.conv2_stride,
                padding=self.dilation,
                dilation=self.dilation,
                groups=groups,
                bias=False)
67

68
        self.add_module(self.norm2_name, norm2)
69
70
71
72
73
74
        self.conv3 = build_conv_layer(
            self.conv_cfg,
            width,
            self.planes * self.expansion,
            kernel_size=1,
            bias=False)
75
        self.add_module(self.norm3_name, norm3)
pangjm's avatar
pangjm committed
76
77
78
79
80
81
82
83
84
85
86


def make_res_layer(block,
                   inplanes,
                   planes,
                   blocks,
                   stride=1,
                   dilation=1,
                   groups=1,
                   base_width=4,
                   style='pytorch',
ThangVu's avatar
ThangVu committed
87
                   with_cp=False,
88
                   conv_cfg=None,
Kai Chen's avatar
Kai Chen committed
89
                   norm_cfg=dict(type='BN'),
90
91
                   dcn=None,
                   gcb=None):
pangjm's avatar
pangjm committed
92
93
94
    downsample = None
    if stride != 1 or inplanes != planes * block.expansion:
        downsample = nn.Sequential(
95
96
            build_conv_layer(
                conv_cfg,
pangjm's avatar
pangjm committed
97
98
99
100
101
                inplanes,
                planes * block.expansion,
                kernel_size=1,
                stride=stride,
                bias=False),
Kai Chen's avatar
Kai Chen committed
102
            build_norm_layer(norm_cfg, planes * block.expansion)[1],
pangjm's avatar
pangjm committed
103
104
105
106
107
        )

    layers = []
    layers.append(
        block(
Kai Chen's avatar
Kai Chen committed
108
109
            inplanes=inplanes,
            planes=planes,
pangjm's avatar
pangjm committed
110
111
112
            stride=stride,
            dilation=dilation,
            downsample=downsample,
pangjm's avatar
pangjm committed
113
114
115
            groups=groups,
            base_width=base_width,
            style=style,
ThangVu's avatar
ThangVu committed
116
            with_cp=with_cp,
117
            conv_cfg=conv_cfg,
Kai Chen's avatar
Kai Chen committed
118
            norm_cfg=norm_cfg,
119
120
            dcn=dcn,
            gcb=gcb))
pangjm's avatar
pangjm committed
121
122
123
124
    inplanes = planes * block.expansion
    for i in range(1, blocks):
        layers.append(
            block(
Kai Chen's avatar
Kai Chen committed
125
126
                inplanes=inplanes,
                planes=planes,
pangjm's avatar
pangjm committed
127
128
                stride=1,
                dilation=dilation,
pangjm's avatar
pangjm committed
129
130
131
                groups=groups,
                base_width=base_width,
                style=style,
ThangVu's avatar
ThangVu committed
132
                with_cp=with_cp,
133
                conv_cfg=conv_cfg,
Kai Chen's avatar
Kai Chen committed
134
                norm_cfg=norm_cfg,
135
136
                dcn=dcn,
                gcb=gcb))
pangjm's avatar
pangjm committed
137
138
139
140

    return nn.Sequential(*layers)


Kai Chen's avatar
Kai Chen committed
141
@BACKBONES.register_module
pangjm's avatar
pangjm committed
142
143
144
145
146
class ResNeXt(ResNet):
    """ResNeXt backbone.

    Args:
        depth (int): Depth of resnet, from {18, 34, 50, 101, 152}.
147
        in_channels (int): Number of input image channels. Normally 3.
pangjm's avatar
pangjm committed
148
149
150
151
152
153
154
155
156
157
158
        num_stages (int): Resnet stages, normally 4.
        groups (int): Group of resnext.
        base_width (int): Base width of resnext.
        strides (Sequence[int]): Strides of the first block of each stage.
        dilations (Sequence[int]): Dilation of each stage.
        out_indices (Sequence[int]): Output from which stages.
        style (str): `pytorch` or `caffe`. If set to "pytorch", the stride-two
            layer is the 3x3 conv layer, otherwise the stride-two layer is
            the first 1x1 conv layer.
        frozen_stages (int): Stages to be frozen (all param fixed). -1 means
            not freezing any parameters.
Kai Chen's avatar
Kai Chen committed
159
        norm_cfg (dict): dictionary to construct and config norm layer.
thangvu's avatar
thangvu committed
160
161
162
        norm_eval (bool): Whether to set norm layers to eval mode, namely,
            freeze running stats (mean and var). Note: Effect on Batch Norm
            and its variants only.
pangjm's avatar
pangjm committed
163
164
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed.
thangvu's avatar
thangvu committed
165
166
        zero_init_residual (bool): whether to use zero init for last norm layer
            in resblocks to let them behave as identity.
167
168
169
170
171
172
173
174
175
176
177
178
179
180

    Example:
        >>> from mmdet.models import ResNeXt
        >>> import torch
        >>> self = ResNeXt(depth=50)
        >>> self.eval()
        >>> inputs = torch.rand(1, 3, 32, 32)
        >>> level_outputs = self.forward(inputs)
        >>> for level_out in level_outputs:
        ...     print(tuple(level_out.shape))
        (1, 256, 8, 8)
        (1, 512, 4, 4)
        (1, 1024, 2, 2)
        (1, 2048, 1, 1)
pangjm's avatar
pangjm committed
181
182
183
184
185
186
187
188
    """

    arch_settings = {
        50: (Bottleneck, (3, 4, 6, 3)),
        101: (Bottleneck, (3, 4, 23, 3)),
        152: (Bottleneck, (3, 8, 36, 3))
    }

pangjm's avatar
pangjm committed
189
190
    def __init__(self, groups=1, base_width=4, **kwargs):
        super(ResNeXt, self).__init__(**kwargs)
pangjm's avatar
pangjm committed
191
192
193
194
195
196
        self.groups = groups
        self.base_width = base_width

        self.inplanes = 64
        self.res_layers = []
        for i, num_blocks in enumerate(self.stage_blocks):
pangjm's avatar
pangjm committed
197
198
            stride = self.strides[i]
            dilation = self.dilations[i]
yhcao6's avatar
yhcao6 committed
199
            dcn = self.dcn if self.stage_with_dcn[i] else None
200
            gcb = self.gcb if self.stage_with_gcb[i] else None
pangjm's avatar
pangjm committed
201
202
203
204
205
206
207
208
209
210
211
            planes = 64 * 2**i
            res_layer = make_res_layer(
                self.block,
                self.inplanes,
                planes,
                num_blocks,
                stride=stride,
                dilation=dilation,
                groups=self.groups,
                base_width=self.base_width,
                style=self.style,
ThangVu's avatar
ThangVu committed
212
                with_cp=self.with_cp,
213
                conv_cfg=self.conv_cfg,
Kai Chen's avatar
Kai Chen committed
214
                norm_cfg=self.norm_cfg,
215
216
                dcn=dcn,
                gcb=gcb)
pangjm's avatar
pangjm committed
217
218
219
220
            self.inplanes = planes * self.block.expansion
            layer_name = 'layer{}'.format(i + 1)
            self.add_module(layer_name, res_layer)
            self.res_layers.append(layer_name)
ThangVu's avatar
ThangVu committed
221
222

        self._freeze_stages()