resnet.py 15 KB
Newer Older
1
2
import logging

Kai Chen's avatar
Kai Chen committed
3
4
import torch.nn as nn
import torch.utils.checkpoint as cp
Kai Chen's avatar
Kai Chen committed
5
6

from mmcv.cnn import constant_init, kaiming_init
Kai Chen's avatar
Kai Chen committed
7
from mmcv.runner import load_checkpoint
Kai Chen's avatar
Kai Chen committed
8
9
10

from mmdet.ops import DeformConv, ModulatedDeformConv
from ..registry import BACKBONES
11
from ..utils import build_conv_layer, build_norm_layer
Kai Chen's avatar
Kai Chen committed
12
13
14
15
16
17
18
19
20
21
22


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 dilation=1,
                 downsample=None,
Kai Chen's avatar
Kai Chen committed
23
                 style='pytorch',
24
                 with_cp=False,
25
                 conv_cfg=None,
Kai Chen's avatar
Kai Chen committed
26
                 norm_cfg=dict(type='BN'),
pangjm's avatar
pangjm committed
27
                 dcn=None):
Kai Chen's avatar
Kai Chen committed
28
        super(BasicBlock, self).__init__()
pangjm's avatar
pangjm committed
29
        assert dcn is None, "Not implemented yet."
30

Kai Chen's avatar
Kai Chen committed
31
32
        self.norm1_name, norm1 = build_norm_layer(norm_cfg, planes, postfix=1)
        self.norm2_name, norm2 = build_norm_layer(norm_cfg, planes, postfix=2)
33

34
35
36
37
38
39
40
41
42
        self.conv1 = build_conv_layer(
            conv_cfg,
            inplanes,
            planes,
            3,
            stride=stride,
            padding=dilation,
            dilation=dilation,
            bias=False)
ThangVu's avatar
ThangVu committed
43
        self.add_module(self.norm1_name, norm1)
44
45
46
47
48
        self.conv2 = build_conv_layer(
            conv_cfg,
            planes,
            planes,
            3,
liuzili97's avatar
liuzili97 committed
49
            padding=1,
50
            bias=False)
ThangVu's avatar
ThangVu committed
51
        self.add_module(self.norm2_name, norm2)
52

Kai Chen's avatar
Kai Chen committed
53
54
55
56
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride
        self.dilation = dilation
Kai Chen's avatar
Kai Chen committed
57
        assert not with_cp
Kai Chen's avatar
Kai Chen committed
58

ThangVu's avatar
ThangVu committed
59
60
61
62
63
64
65
66
    @property
    def norm1(self):
        return getattr(self, self.norm1_name)

    @property
    def norm2(self):
        return getattr(self, self.norm2_name)

Kai Chen's avatar
Kai Chen committed
67
    def forward(self, x):
pangjm's avatar
pangjm committed
68
        identity = x
Kai Chen's avatar
Kai Chen committed
69
70

        out = self.conv1(x)
ThangVu's avatar
ThangVu committed
71
        out = self.norm1(out)
Kai Chen's avatar
Kai Chen committed
72
73
74
        out = self.relu(out)

        out = self.conv2(out)
ThangVu's avatar
ThangVu committed
75
        out = self.norm2(out)
Kai Chen's avatar
Kai Chen committed
76
77

        if self.downsample is not None:
pangjm's avatar
pangjm committed
78
            identity = self.downsample(x)
Kai Chen's avatar
Kai Chen committed
79

pangjm's avatar
pangjm committed
80
        out += identity
Kai Chen's avatar
Kai Chen committed
81
82
83
84
85
86
87
88
89
90
91
92
93
94
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 dilation=1,
                 downsample=None,
95
                 style='pytorch',
96
                 with_cp=False,
97
                 conv_cfg=None,
Kai Chen's avatar
Kai Chen committed
98
                 norm_cfg=dict(type='BN'),
Kai Chen's avatar
Kai Chen committed
99
                 dcn=None):
pangjm's avatar
pangjm committed
100
        """Bottleneck block for ResNet.
101
102
        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
103
104
        """
        super(Bottleneck, self).__init__()
105
        assert style in ['pytorch', 'caffe']
Kai Chen's avatar
Kai Chen committed
106
        assert dcn is None or isinstance(dcn, dict)
pangjm's avatar
pangjm committed
107
108
        self.inplanes = inplanes
        self.planes = planes
Kai Chen's avatar
Kai Chen committed
109
110
111
112
        self.stride = stride
        self.dilation = dilation
        self.style = style
        self.with_cp = with_cp
113
        self.conv_cfg = conv_cfg
Kai Chen's avatar
Kai Chen committed
114
        self.norm_cfg = norm_cfg
Kai Chen's avatar
Kai Chen committed
115
116
        self.dcn = dcn
        self.with_dcn = dcn is not None
Kai Chen's avatar
Kai Chen committed
117
        if self.style == 'pytorch':
pangjm's avatar
pangjm committed
118
119
            self.conv1_stride = 1
            self.conv2_stride = stride
Kai Chen's avatar
Kai Chen committed
120
        else:
pangjm's avatar
pangjm committed
121
122
            self.conv1_stride = stride
            self.conv2_stride = 1
123

Kai Chen's avatar
Kai Chen committed
124
125
        self.norm1_name, norm1 = build_norm_layer(norm_cfg, planes, postfix=1)
        self.norm2_name, norm2 = build_norm_layer(norm_cfg, planes, postfix=2)
Kai Chen's avatar
Kai Chen committed
126
        self.norm3_name, norm3 = build_norm_layer(
Kai Chen's avatar
Kai Chen committed
127
            norm_cfg, planes * self.expansion, postfix=3)
128

129
130
        self.conv1 = build_conv_layer(
            conv_cfg,
pangjm's avatar
pangjm committed
131
132
133
134
135
            inplanes,
            planes,
            kernel_size=1,
            stride=self.conv1_stride,
            bias=False)
136
        self.add_module(self.norm1_name, norm1)
Kai Chen's avatar
Kai Chen committed
137
138
139
140
141
142
        fallback_on_stride = False
        self.with_modulated_dcn = False
        if self.with_dcn:
            fallback_on_stride = dcn.get('fallback_on_stride', False)
            self.with_modulated_dcn = dcn.get('modulated', False)
        if not self.with_dcn or fallback_on_stride:
143
144
            self.conv2 = build_conv_layer(
                conv_cfg,
Kai Chen's avatar
Kai Chen committed
145
146
147
148
149
150
151
152
                planes,
                planes,
                kernel_size=3,
                stride=self.conv2_stride,
                padding=dilation,
                dilation=dilation,
                bias=False)
        else:
153
            assert conv_cfg is None, 'conv_cfg must be None for DCN'
Kai Chen's avatar
Kai Chen committed
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
            deformable_groups = dcn.get('deformable_groups', 1)
            if not self.with_modulated_dcn:
                conv_op = DeformConv
                offset_channels = 18
            else:
                conv_op = ModulatedDeformConv
                offset_channels = 27
            self.conv2_offset = nn.Conv2d(
                planes,
                deformable_groups * offset_channels,
                kernel_size=3,
                stride=self.conv2_stride,
                padding=dilation,
                dilation=dilation)
            self.conv2 = conv_op(
                planes,
                planes,
                kernel_size=3,
                stride=self.conv2_stride,
                padding=dilation,
                dilation=dilation,
                deformable_groups=deformable_groups,
                bias=False)
ThangVu's avatar
ThangVu committed
177
        self.add_module(self.norm2_name, norm2)
178
179
180
181
182
183
        self.conv3 = build_conv_layer(
            conv_cfg,
            planes,
            planes * self.expansion,
            kernel_size=1,
            bias=False)
184
185
        self.add_module(self.norm3_name, norm3)

Kai Chen's avatar
Kai Chen committed
186
        self.relu = nn.ReLU(inplace=True)
187
        self.downsample = downsample
Kai Chen's avatar
Kai Chen committed
188

ThangVu's avatar
ThangVu committed
189
190
191
192
193
194
195
196
197
198
199
200
    @property
    def norm1(self):
        return getattr(self, self.norm1_name)

    @property
    def norm2(self):
        return getattr(self, self.norm2_name)

    @property
    def norm3(self):
        return getattr(self, self.norm3_name)

Kai Chen's avatar
Kai Chen committed
201
202
203
    def forward(self, x):

        def _inner_forward(x):
pangjm's avatar
pangjm committed
204
            identity = x
Kai Chen's avatar
Kai Chen committed
205
206

            out = self.conv1(x)
ThangVu's avatar
ThangVu committed
207
            out = self.norm1(out)
Kai Chen's avatar
Kai Chen committed
208
209
            out = self.relu(out)

Kai Chen's avatar
Kai Chen committed
210
211
212
213
214
215
216
217
218
219
            if not self.with_dcn:
                out = self.conv2(out)
            elif self.with_modulated_dcn:
                offset_mask = self.conv2_offset(out)
                offset = offset_mask[:, :18, :, :]
                mask = offset_mask[:, -9:, :, :].sigmoid()
                out = self.conv2(out, offset, mask)
            else:
                offset = self.conv2_offset(out)
                out = self.conv2(out, offset)
ThangVu's avatar
ThangVu committed
220
            out = self.norm2(out)
Kai Chen's avatar
Kai Chen committed
221
222
223
            out = self.relu(out)

            out = self.conv3(out)
ThangVu's avatar
ThangVu committed
224
            out = self.norm3(out)
Kai Chen's avatar
Kai Chen committed
225
226

            if self.downsample is not None:
pangjm's avatar
pangjm committed
227
                identity = self.downsample(x)
Kai Chen's avatar
Kai Chen committed
228

pangjm's avatar
pangjm committed
229
            out += identity
Kai Chen's avatar
Kai Chen committed
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248

            return out

        if self.with_cp and x.requires_grad:
            out = cp.checkpoint(_inner_forward, x)
        else:
            out = _inner_forward(x)

        out = self.relu(out)

        return out


def make_res_layer(block,
                   inplanes,
                   planes,
                   blocks,
                   stride=1,
                   dilation=1,
249
                   style='pytorch',
250
                   with_cp=False,
251
                   conv_cfg=None,
Kai Chen's avatar
Kai Chen committed
252
                   norm_cfg=dict(type='BN'),
Kai Chen's avatar
Kai Chen committed
253
                   dcn=None):
Kai Chen's avatar
Kai Chen committed
254
255
256
    downsample = None
    if stride != 1 or inplanes != planes * block.expansion:
        downsample = nn.Sequential(
257
258
            build_conv_layer(
                conv_cfg,
Kai Chen's avatar
Kai Chen committed
259
260
261
262
263
                inplanes,
                planes * block.expansion,
                kernel_size=1,
                stride=stride,
                bias=False),
Kai Chen's avatar
Kai Chen committed
264
            build_norm_layer(norm_cfg, planes * block.expansion)[1],
Kai Chen's avatar
Kai Chen committed
265
266
267
268
269
270
271
272
273
274
275
        )

    layers = []
    layers.append(
        block(
            inplanes,
            planes,
            stride,
            dilation,
            downsample,
            style=style,
276
            with_cp=with_cp,
277
            conv_cfg=conv_cfg,
Kai Chen's avatar
Kai Chen committed
278
            norm_cfg=norm_cfg,
Kai Chen's avatar
Kai Chen committed
279
            dcn=dcn))
Kai Chen's avatar
Kai Chen committed
280
281
282
    inplanes = planes * block.expansion
    for i in range(1, blocks):
        layers.append(
Kai Chen's avatar
Kai Chen committed
283
284
285
286
287
288
289
            block(
                inplanes,
                planes,
                1,
                dilation,
                style=style,
                with_cp=with_cp,
290
                conv_cfg=conv_cfg,
Kai Chen's avatar
Kai Chen committed
291
                norm_cfg=norm_cfg,
Kai Chen's avatar
Kai Chen committed
292
                dcn=dcn))
Kai Chen's avatar
Kai Chen committed
293
294
295
296

    return nn.Sequential(*layers)


Kai Chen's avatar
Kai Chen committed
297
@BACKBONES.register_module
Kai Chen's avatar
Kai Chen committed
298
299
class ResNet(nn.Module):
    """ResNet backbone.
Kai Chen's avatar
Kai Chen committed
300

Kai Chen's avatar
Kai Chen committed
301
302
303
304
305
306
307
308
309
    Args:
        depth (int): Depth of resnet, from {18, 34, 50, 101, 152}.
        num_stages (int): Resnet stages, normally 4.
        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.
310
311
        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
            -1 means not freezing any parameters.
Kai Chen's avatar
Kai Chen committed
312
        norm_cfg (dict): dictionary to construct and config norm layer.
thangvu's avatar
thangvu committed
313
314
315
        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.
Kai Chen's avatar
Kai Chen committed
316
317
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed.
thangvu's avatar
thangvu committed
318
319
        zero_init_residual (bool): whether to use zero init for last norm layer
            in resblocks to let them behave as identity.
Kai Chen's avatar
Kai Chen committed
320
    """
Kai Chen's avatar
Kai Chen committed
321

Kai Chen's avatar
Kai Chen committed
322
323
324
325
326
327
328
    arch_settings = {
        18: (BasicBlock, (2, 2, 2, 2)),
        34: (BasicBlock, (3, 4, 6, 3)),
        50: (Bottleneck, (3, 4, 6, 3)),
        101: (Bottleneck, (3, 4, 23, 3)),
        152: (Bottleneck, (3, 8, 36, 3))
    }
Kai Chen's avatar
Kai Chen committed
329
330

    def __init__(self,
Kai Chen's avatar
Kai Chen committed
331
332
                 depth,
                 num_stages=4,
Kai Chen's avatar
Kai Chen committed
333
334
335
                 strides=(1, 2, 2, 2),
                 dilations=(1, 1, 1, 1),
                 out_indices=(0, 1, 2, 3),
336
                 style='pytorch',
ThangVu's avatar
ThangVu committed
337
                 frozen_stages=-1,
338
                 conv_cfg=None,
Kai Chen's avatar
Kai Chen committed
339
                 norm_cfg=dict(type='BN', requires_grad=True),
thangvu's avatar
thangvu committed
340
                 norm_eval=True,
Kai Chen's avatar
Kai Chen committed
341
342
                 dcn=None,
                 stage_with_dcn=(False, False, False, False),
ThangVu's avatar
ThangVu committed
343
344
                 with_cp=False,
                 zero_init_residual=True):
Kai Chen's avatar
Kai Chen committed
345
        super(ResNet, self).__init__()
Kai Chen's avatar
Kai Chen committed
346
347
        if depth not in self.arch_settings:
            raise KeyError('invalid depth {} for resnet'.format(depth))
pangjm's avatar
pangjm committed
348
349
        self.depth = depth
        self.num_stages = num_stages
Kai Chen's avatar
Kai Chen committed
350
        assert num_stages >= 1 and num_stages <= 4
pangjm's avatar
pangjm committed
351
352
        self.strides = strides
        self.dilations = dilations
Kai Chen's avatar
Kai Chen committed
353
        assert len(strides) == len(dilations) == num_stages
Kai Chen's avatar
Kai Chen committed
354
        self.out_indices = out_indices
pangjm's avatar
pangjm committed
355
        assert max(out_indices) < num_stages
Kai Chen's avatar
Kai Chen committed
356
        self.style = style
ThangVu's avatar
ThangVu committed
357
        self.frozen_stages = frozen_stages
358
        self.conv_cfg = conv_cfg
Kai Chen's avatar
Kai Chen committed
359
        self.norm_cfg = norm_cfg
ThangVu's avatar
ThangVu committed
360
        self.with_cp = with_cp
thangvu's avatar
thangvu committed
361
        self.norm_eval = norm_eval
Kai Chen's avatar
Kai Chen committed
362
363
        self.dcn = dcn
        self.stage_with_dcn = stage_with_dcn
Kai Chen's avatar
Kai Chen committed
364
365
        if dcn is not None:
            assert len(stage_with_dcn) == num_stages
ThangVu's avatar
ThangVu committed
366
        self.zero_init_residual = zero_init_residual
pangjm's avatar
pangjm committed
367
368
        self.block, stage_blocks = self.arch_settings[depth]
        self.stage_blocks = stage_blocks[:num_stages]
Kai Chen's avatar
Kai Chen committed
369
        self.inplanes = 64
pangjm's avatar
pangjm committed
370

thangvu's avatar
thangvu committed
371
        self._make_stem_layer()
Kai Chen's avatar
Kai Chen committed
372

Kai Chen's avatar
Kai Chen committed
373
        self.res_layers = []
pangjm's avatar
pangjm committed
374
        for i, num_blocks in enumerate(self.stage_blocks):
Kai Chen's avatar
Kai Chen committed
375
376
            stride = strides[i]
            dilation = dilations[i]
Kai Chen's avatar
Kai Chen committed
377
            dcn = self.dcn if self.stage_with_dcn[i] else None
Kai Chen's avatar
Kai Chen committed
378
379
            planes = 64 * 2**i
            res_layer = make_res_layer(
pangjm's avatar
pangjm committed
380
                self.block,
Kai Chen's avatar
Kai Chen committed
381
382
383
384
385
386
                self.inplanes,
                planes,
                num_blocks,
                stride=stride,
                dilation=dilation,
                style=self.style,
387
                with_cp=with_cp,
388
                conv_cfg=conv_cfg,
Kai Chen's avatar
Kai Chen committed
389
                norm_cfg=norm_cfg,
Kai Chen's avatar
Kai Chen committed
390
                dcn=dcn)
pangjm's avatar
pangjm committed
391
            self.inplanes = planes * self.block.expansion
Kai Chen's avatar
Kai Chen committed
392
            layer_name = 'layer{}'.format(i + 1)
393
            self.add_module(layer_name, res_layer)
Kai Chen's avatar
Kai Chen committed
394
395
            self.res_layers.append(layer_name)

ThangVu's avatar
ThangVu committed
396
397
        self._freeze_stages()

pangjm's avatar
pangjm committed
398
399
        self.feat_dim = self.block.expansion * 64 * 2**(
            len(self.stage_blocks) - 1)
pangjm's avatar
pangjm committed
400

ThangVu's avatar
ThangVu committed
401
402
403
404
    @property
    def norm1(self):
        return getattr(self, self.norm1_name)

thangvu's avatar
thangvu committed
405
    def _make_stem_layer(self):
406
407
408
409
410
411
412
413
        self.conv1 = build_conv_layer(
            self.conv_cfg,
            3,
            64,
            kernel_size=7,
            stride=2,
            padding=3,
            bias=False)
Kai Chen's avatar
Kai Chen committed
414
        self.norm1_name, norm1 = build_norm_layer(self.norm_cfg, 64, postfix=1)
ThangVu's avatar
ThangVu committed
415
        self.add_module(self.norm1_name, norm1)
thangvu's avatar
thangvu committed
416
417
418
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

ThangVu's avatar
ThangVu committed
419
    def _freeze_stages(self):
ThangVu's avatar
ThangVu committed
420
        if self.frozen_stages >= 0:
Kai Chen's avatar
Kai Chen committed
421
            self.norm1.eval()
ThangVu's avatar
ThangVu committed
422
            for m in [self.conv1, self.norm1]:
ThangVu's avatar
ThangVu committed
423
                for param in m.parameters():
thangvu's avatar
thangvu committed
424
425
                    param.requires_grad = False

ThangVu's avatar
ThangVu committed
426
427
        for i in range(1, self.frozen_stages + 1):
            m = getattr(self, 'layer{}'.format(i))
Kai Chen's avatar
Kai Chen committed
428
            m.eval()
ThangVu's avatar
ThangVu committed
429
430
431
            for param in m.parameters():
                param.requires_grad = False

Kai Chen's avatar
Kai Chen committed
432
433
    def init_weights(self, pretrained=None):
        if isinstance(pretrained, str):
434
435
            logger = logging.getLogger()
            load_checkpoint(self, pretrained, strict=False, logger=logger)
Kai Chen's avatar
Kai Chen committed
436
437
438
        elif pretrained is None:
            for m in self.modules():
                if isinstance(m, nn.Conv2d):
Kai Chen's avatar
Kai Chen committed
439
                    kaiming_init(m)
ThangVu's avatar
minor  
ThangVu committed
440
                elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
Kai Chen's avatar
Kai Chen committed
441
                    constant_init(m, 1)
442

Kai Chen's avatar
Kai Chen committed
443
444
445
446
447
448
            if self.dcn is not None:
                for m in self.modules():
                    if isinstance(m, Bottleneck) and hasattr(
                            m, 'conv2_offset'):
                        constant_init(m.conv2_offset, 0)

ThangVu's avatar
ThangVu committed
449
450
            if self.zero_init_residual:
                for m in self.modules():
ThangVu's avatar
ThangVu committed
451
452
453
454
                    if isinstance(m, Bottleneck):
                        constant_init(m.norm3, 0)
                    elif isinstance(m, BasicBlock):
                        constant_init(m.norm2, 0)
Kai Chen's avatar
Kai Chen committed
455
456
457
458
459
        else:
            raise TypeError('pretrained must be a str or None')

    def forward(self, x):
        x = self.conv1(x)
ThangVu's avatar
ThangVu committed
460
        x = self.norm1(x)
Kai Chen's avatar
Kai Chen committed
461
462
463
464
465
466
467
468
        x = self.relu(x)
        x = self.maxpool(x)
        outs = []
        for i, layer_name in enumerate(self.res_layers):
            res_layer = getattr(self, layer_name)
            x = res_layer(x)
            if i in self.out_indices:
                outs.append(x)
myownskyW7's avatar
myownskyW7 committed
469
        return tuple(outs)
Kai Chen's avatar
Kai Chen committed
470
471
472

    def train(self, mode=True):
        super(ResNet, self).train(mode)
thangvu's avatar
thangvu committed
473
        if mode and self.norm_eval:
ThangVu's avatar
ThangVu committed
474
            for m in self.modules():
thangvu's avatar
thangvu committed
475
                # trick: eval have effect on BatchNorm only
ThangVu's avatar
ThangVu committed
476
477
                if isinstance(m, nn.BatchNorm2d):
                    m.eval()