Arulkumar03/Fox_Sheep_Detector_Computer_Vision_model
0
1# Copyright (c) Facebook, Inc. and its affiliates.2import numpy as np3import fvcore.nn.weight_init as weight_init4import torch5import torch.nn.functional as F6from torch import nn7 8from detectron2.layers import (9 CNNBlockBase,10 Conv2d,11 DeformConv,12 ModulatedDeformConv,13 ShapeSpec,14 get_norm,15)16 17from .backbone import Backbone18from .build import BACKBONE_REGISTRY19 20__all__ = [21 "ResNetBlockBase",22 "BasicBlock",23 "BottleneckBlock",24 "DeformBottleneckBlock",25 "BasicStem",26 "ResNet",27 "make_stage",28 "build_resnet_backbone",29]30 31 32class BasicBlock(CNNBlockBase):33 """34 The basic residual block for ResNet-18 and ResNet-34 defined in :paper:`ResNet`,35 with two 3x3 conv layers and a projection shortcut if needed.36 """37 38 def __init__(self, in_channels, out_channels, *, stride=1, norm="BN"):39 """40 Args:41 in_channels (int): Number of input channels.42 out_channels (int): Number of output channels.43 stride (int): Stride for the first conv.44 norm (str or callable): normalization for all conv layers.45 See :func:`layers.get_norm` for supported format.46 """47 super().__init__(in_channels, out_channels, stride)48 49 if in_channels != out_channels:50 self.shortcut = Conv2d(51 in_channels,52 out_channels,53 kernel_size=1,54 stride=stride,55 bias=False,56 norm=get_norm(norm, out_channels),57 )58 else:59 self.shortcut = None60 61 self.conv1 = Conv2d(62 in_channels,63 out_channels,64 kernel_size=3,65 stride=stride,66 padding=1,67 bias=False,68 norm=get_norm(norm, out_channels),69 )70 71 self.conv2 = Conv2d(72 out_channels,73 out_channels,74 kernel_size=3,75 stride=1,76 padding=1,77 bias=False,78 norm=get_norm(norm, out_channels),79 )80 81 for layer in [self.conv1, self.conv2, self.shortcut]:82 if layer is not None: # shortcut can be None83 weight_init.c2_msra_fill(layer)84 85 def forward(self, x):86 out = self.conv1(x)87 out = F.relu_(out)88 out = self.conv2(out)89 90 if self.shortcut is not None:91 shortcut = self.shortcut(x)92 else:93 shortcut = x94 95 out += shortcut96 out = F.relu_(out)97 return out98 99 100class BottleneckBlock(CNNBlockBase):101 """102 The standard bottleneck residual block used by ResNet-50, 101 and 152103 defined in :paper:`ResNet`. It contains 3 conv layers with kernels104 1x1, 3x3, 1x1, and a projection shortcut if needed.105 """106 107 def __init__(108 self,109 in_channels,110 out_channels,111 *,112 bottleneck_channels,113 stride=1,114 num_groups=1,115 norm="BN",116 stride_in_1x1=False,117 dilation=1,118 ):119 """120 Args:121 bottleneck_channels (int): number of output channels for the 3x3122 "bottleneck" conv layers.123 num_groups (int): number of groups for the 3x3 conv layer.124 norm (str or callable): normalization for all conv layers.125 See :func:`layers.get_norm` for supported format.126 stride_in_1x1 (bool): when stride>1, whether to put stride in the127 first 1x1 convolution or the bottleneck 3x3 convolution.128 dilation (int): the dilation rate of the 3x3 conv layer.129 """130 super().__init__(in_channels, out_channels, stride)131 132 if in_channels != out_channels:133 self.shortcut = Conv2d(134 in_channels,135 out_channels,136 kernel_size=1,137 stride=stride,138 bias=False,139 norm=get_norm(norm, out_channels),140 )141 else:142 self.shortcut = None143 144 # The original MSRA ResNet models have stride in the first 1x1 conv145 # The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have146 # stride in the 3x3 conv147 stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)148 149 self.conv1 = Conv2d(150 in_channels,151 bottleneck_channels,152 kernel_size=1,153 stride=stride_1x1,154 bias=False,155 norm=get_norm(norm, bottleneck_channels),156 )157 158 self.conv2 = Conv2d(159 bottleneck_channels,160 bottleneck_channels,161 kernel_size=3,162 stride=stride_3x3,163 padding=1 * dilation,164 bias=False,165 groups=num_groups,166 dilation=dilation,167 norm=get_norm(norm, bottleneck_channels),168 )169 170 self.conv3 = Conv2d(171 bottleneck_channels,172 out_channels,173 kernel_size=1,174 bias=False,175 norm=get_norm(norm, out_channels),176 )177 178 for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:179 if layer is not None: # shortcut can be None180 weight_init.c2_msra_fill(layer)181 182 # Zero-initialize the last normalization in each residual branch,183 # so that at the beginning, the residual branch starts with zeros,184 # and each residual block behaves like an identity.185 # See Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":186 # "For BN layers, the learnable scaling coefficient γ is initialized187 # to be 1, except for each residual block's last BN188 # where γ is initialized to be 0."189 190 # nn.init.constant_(self.conv3.norm.weight, 0)191 # TODO this somehow hurts performance when training GN models from scratch.192 # Add it as an option when we need to use this code to train a backbone.193 194 def forward(self, x):195 out = self.conv1(x)196 out = F.relu_(out)197 198 out = self.conv2(out)199 out = F.relu_(out)200 201 out = self.conv3(out)202 203 if self.shortcut is not None:204 shortcut = self.shortcut(x)205 else:206 shortcut = x207 208 out += shortcut209 out = F.relu_(out)210 return out211 212 213class DeformBottleneckBlock(CNNBlockBase):214 """215 Similar to :class:`BottleneckBlock`, but with :paper:`deformable conv <deformconv>`216 in the 3x3 convolution.217 """218 219 def __init__(220 self,221 in_channels,222 out_channels,223 *,224 bottleneck_channels,225 stride=1,226 num_groups=1,227 norm="BN",228 stride_in_1x1=False,229 dilation=1,230 deform_modulated=False,231 deform_num_groups=1,232 ):233 super().__init__(in_channels, out_channels, stride)234 self.deform_modulated = deform_modulated235 236 if in_channels != out_channels:237 self.shortcut = Conv2d(238 in_channels,239 out_channels,240 kernel_size=1,241 stride=stride,242 bias=False,243 norm=get_norm(norm, out_channels),244 )245 else:246 self.shortcut = None247 248 stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)249 250 self.conv1 = Conv2d(251 in_channels,252 bottleneck_channels,253 kernel_size=1,254 stride=stride_1x1,255 bias=False,256 norm=get_norm(norm, bottleneck_channels),257 )258 259 if deform_modulated:260 deform_conv_op = ModulatedDeformConv261 # offset channels are 2 or 3 (if with modulated) * kernel_size * kernel_size262 offset_channels = 27263 else:264 deform_conv_op = DeformConv265 offset_channels = 18266 267 self.conv2_offset = Conv2d(268 bottleneck_channels,269 offset_channels * deform_num_groups,270 kernel_size=3,271 stride=stride_3x3,272 padding=1 * dilation,273 dilation=dilation,274 )275 self.conv2 = deform_conv_op(276 bottleneck_channels,277 bottleneck_channels,278 kernel_size=3,279 stride=stride_3x3,280 padding=1 * dilation,281 bias=False,282 groups=num_groups,283 dilation=dilation,284 deformable_groups=deform_num_groups,285 norm=get_norm(norm, bottleneck_channels),286 )287 288 self.conv3 = Conv2d(289 bottleneck_channels,290 out_channels,291 kernel_size=1,292 bias=False,293 norm=get_norm(norm, out_channels),294 )295 296 for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:297 if layer is not None: # shortcut can be None298 weight_init.c2_msra_fill(layer)299 300 nn.init.constant_(self.conv2_offset.weight, 0)301 nn.init.constant_(self.conv2_offset.bias, 0)302 303 def forward(self, x):304 out = self.conv1(x)305 out = F.relu_(out)306 307 if self.deform_modulated:308 offset_mask = self.conv2_offset(out)309 offset_x, offset_y, mask = torch.chunk(offset_mask, 3, dim=1)310 offset = torch.cat((offset_x, offset_y), dim=1)311 mask = mask.sigmoid()312 out = self.conv2(out, offset, mask)313 else:314 offset = self.conv2_offset(out)315 out = self.conv2(out, offset)316 out = F.relu_(out)317 318 out = self.conv3(out)319 320 if self.shortcut is not None:321 shortcut = self.shortcut(x)322 else:323 shortcut = x324 325 out += shortcut326 out = F.relu_(out)327 return out328 329 330class BasicStem(CNNBlockBase):331 """332 The standard ResNet stem (layers before the first residual block),333 with a conv, relu and max_pool.334 """335 336 def __init__(self, in_channels=3, out_channels=64, norm="BN"):337 """338 Args:339 norm (str or callable): norm after the first conv layer.340 See :func:`layers.get_norm` for supported format.341 """342 super().__init__(in_channels, out_channels, 4)343 self.in_channels = in_channels344 self.conv1 = Conv2d(345 in_channels,346 out_channels,347 kernel_size=7,348 stride=2,349 padding=3,350 bias=False,351 norm=get_norm(norm, out_channels),352 )353 weight_init.c2_msra_fill(self.conv1)354 355 def forward(self, x):356 x = self.conv1(x)357 x = F.relu_(x)358 x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)359 return x360 361 362class ResNet(Backbone):363 """364 Implement :paper:`ResNet`.365 """366 367 def __init__(self, stem, stages, num_classes=None, out_features=None, freeze_at=0):368 """369 Args:370 stem (nn.Module): a stem module371 stages (list[list[CNNBlockBase]]): several (typically 4) stages,372 each contains multiple :class:`CNNBlockBase`.373 num_classes (None or int): if None, will not perform classification.374 Otherwise, will create a linear layer.375 out_features (list[str]): name of the layers whose outputs should376 be returned in forward. Can be anything in "stem", "linear", or "res2" ...377 If None, will return the output of the last layer.378 freeze_at (int): The number of stages at the beginning to freeze.379 see :meth:`freeze` for detailed explanation.380 """381 super().__init__()382 self.stem = stem383 self.num_classes = num_classes384 385 current_stride = self.stem.stride386 self._out_feature_strides = {"stem": current_stride}387 self._out_feature_channels = {"stem": self.stem.out_channels}388 389 self.stage_names, self.stages = [], []390 391 if out_features is not None:392 # Avoid keeping unused layers in this module. They consume extra memory393 # and may cause allreduce to fail394 num_stages = max(395 [{"res2": 1, "res3": 2, "res4": 3, "res5": 4}.get(f, 0) for f in out_features]396 )397 stages = stages[:num_stages]398 for i, blocks in enumerate(stages):399 assert len(blocks) > 0, len(blocks)400 for block in blocks:401 assert isinstance(block, CNNBlockBase), block402 403 name = "res" + str(i + 2)404 stage = nn.Sequential(*blocks)405 406 self.add_module(name, stage)407 self.stage_names.append(name)408 self.stages.append(stage)409 410 self._out_feature_strides[name] = current_stride = int(411 current_stride * np.prod([k.stride for k in blocks])412 )413 self._out_feature_channels[name] = curr_channels = blocks[-1].out_channels414 self.stage_names = tuple(self.stage_names) # Make it static for scripting415 416 if num_classes is not None:417 self.avgpool = nn.AdaptiveAvgPool2d((1, 1))418 self.linear = nn.Linear(curr_channels, num_classes)419 420 # Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":421 # "The 1000-way fully-connected layer is initialized by422 # drawing weights from a zero-mean Gaussian with standard deviation of 0.01."423 nn.init.normal_(self.linear.weight, std=0.01)424 name = "linear"425 426 if out_features is None:427 out_features = [name]428 self._out_features = out_features429 assert len(self._out_features)430 children = [x[0] for x in self.named_children()]431 for out_feature in self._out_features:432 assert out_feature in children, "Available children: {}".format(", ".join(children))433 self.freeze(freeze_at)434 435 def forward(self, x):436 """437 Args:438 x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.439 440 Returns:441 dict[str->Tensor]: names and the corresponding features442 """443 assert x.dim() == 4, f"ResNet takes an input of shape (N, C, H, W). Got {x.shape} instead!"444 outputs = {}445 x = self.stem(x)446 if "stem" in self._out_features:447 outputs["stem"] = x448 for name, stage in zip(self.stage_names, self.stages):449 x = stage(x)450 if name in self._out_features:451 outputs[name] = x452 if self.num_classes is not None:453 x = self.avgpool(x)454 x = torch.flatten(x, 1)455 x = self.linear(x)456 if "linear" in self._out_features:457 outputs["linear"] = x458 return outputs459 460 def output_shape(self):461 return {462 name: ShapeSpec(463 channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]464 )465 for name in self._out_features466 }467 468 def freeze(self, freeze_at=0):469 """470 Freeze the first several stages of the ResNet. Commonly used in471 fine-tuning.472 473 Layers that produce the same feature map spatial size are defined as one474 "stage" by :paper:`FPN`.475 476 Args:477 freeze_at (int): number of stages to freeze.478 `1` means freezing the stem. `2` means freezing the stem and479 one residual stage, etc.480 481 Returns:482 nn.Module: this ResNet itself483 """484 if freeze_at >= 1:485 self.stem.freeze()486 for idx, stage in enumerate(self.stages, start=2):487 if freeze_at >= idx:488 for block in stage.children():489 block.freeze()490 return self491 492 @staticmethod493 def make_stage(block_class, num_blocks, *, in_channels, out_channels, **kwargs):494 """495 Create a list of blocks of the same type that forms one ResNet stage.496 497 Args:498 block_class (type): a subclass of CNNBlockBase that's used to create all blocks in this499 stage. A module of this type must not change spatial resolution of inputs unless its500 stride != 1.501 num_blocks (int): number of blocks in this stage502 in_channels (int): input channels of the entire stage.503 out_channels (int): output channels of **every block** in the stage.504 kwargs: other arguments passed to the constructor of505 `block_class`. If the argument name is "xx_per_block", the506 argument is a list of values to be passed to each block in the507 stage. Otherwise, the same argument is passed to every block508 in the stage.509 510 Returns:511 list[CNNBlockBase]: a list of block module.512 513 Examples:514 ::515 stage = ResNet.make_stage(516 BottleneckBlock, 3, in_channels=16, out_channels=64,517 bottleneck_channels=16, num_groups=1,518 stride_per_block=[2, 1, 1],519 dilations_per_block=[1, 1, 2]520 )521 522 Usually, layers that produce the same feature map spatial size are defined as one523 "stage" (in :paper:`FPN`). Under such definition, ``stride_per_block[1:]`` should524 all be 1.525 """526 blocks = []527 for i in range(num_blocks):528 curr_kwargs = {}529 for k, v in kwargs.items():530 if k.endswith("_per_block"):531 assert len(v) == num_blocks, (532 f"Argument '{k}' of make_stage should have the "533 f"same length as num_blocks={num_blocks}."534 )535 newk = k[: -len("_per_block")]536 assert newk not in kwargs, f"Cannot call make_stage with both {k} and {newk}!"537 curr_kwargs[newk] = v[i]538 else:539 curr_kwargs[k] = v540 541 blocks.append(542 block_class(in_channels=in_channels, out_channels=out_channels, **curr_kwargs)543 )544 in_channels = out_channels545 return blocks546 547 @staticmethod548 def make_default_stages(depth, block_class=None, **kwargs):549 """550 Created list of ResNet stages from pre-defined depth (one of 18, 34, 50, 101, 152).551 If it doesn't create the ResNet variant you need, please use :meth:`make_stage`552 instead for fine-grained customization.553 554 Args:555 depth (int): depth of ResNet556 block_class (type): the CNN block class. Has to accept557 `bottleneck_channels` argument for depth > 50.558 By default it is BasicBlock or BottleneckBlock, based on the559 depth.560 kwargs:561 other arguments to pass to `make_stage`. Should not contain562 stride and channels, as they are predefined for each depth.563 564 Returns:565 list[list[CNNBlockBase]]: modules in all stages; see arguments of566 :class:`ResNet.__init__`.567 """568 num_blocks_per_stage = {569 18: [2, 2, 2, 2],570 34: [3, 4, 6, 3],571 50: [3, 4, 6, 3],572 101: [3, 4, 23, 3],573 152: [3, 8, 36, 3],574 }[depth]575 if block_class is None:576 block_class = BasicBlock if depth < 50 else BottleneckBlock577 if depth < 50:578 in_channels = [64, 64, 128, 256]579 out_channels = [64, 128, 256, 512]580 else:581 in_channels = [64, 256, 512, 1024]582 out_channels = [256, 512, 1024, 2048]583 ret = []584 for (n, s, i, o) in zip(num_blocks_per_stage, [1, 2, 2, 2], in_channels, out_channels):585 if depth >= 50:586 kwargs["bottleneck_channels"] = o // 4587 ret.append(588 ResNet.make_stage(589 block_class=block_class,590 num_blocks=n,591 stride_per_block=[s] + [1] * (n - 1),592 in_channels=i,593 out_channels=o,594 **kwargs,595 )596 )597 return ret598 599 600ResNetBlockBase = CNNBlockBase601"""602Alias for backward compatibiltiy.603"""604 605 606def make_stage(*args, **kwargs):607 """608 Deprecated alias for backward compatibiltiy.609 """610 return ResNet.make_stage(*args, **kwargs)611 612 613@BACKBONE_REGISTRY.register()614def build_resnet_backbone(cfg, input_shape):615 """616 Create a ResNet instance from config.617 618 Returns:619 ResNet: a :class:`ResNet` instance.620 """621 # need registration of new blocks/stems?622 norm = cfg.MODEL.RESNETS.NORM623 stem = BasicStem(624 in_channels=input_shape.channels,625 out_channels=cfg.MODEL.RESNETS.STEM_OUT_CHANNELS,626 norm=norm,627 )628 629 # fmt: off630 freeze_at = cfg.MODEL.BACKBONE.FREEZE_AT631 out_features = cfg.MODEL.RESNETS.OUT_FEATURES632 depth = cfg.MODEL.RESNETS.DEPTH633 num_groups = cfg.MODEL.RESNETS.NUM_GROUPS634 width_per_group = cfg.MODEL.RESNETS.WIDTH_PER_GROUP635 bottleneck_channels = num_groups * width_per_group636 in_channels = cfg.MODEL.RESNETS.STEM_OUT_CHANNELS637 out_channels = cfg.MODEL.RESNETS.RES2_OUT_CHANNELS638 stride_in_1x1 = cfg.MODEL.RESNETS.STRIDE_IN_1X1639 res5_dilation = cfg.MODEL.RESNETS.RES5_DILATION640 deform_on_per_stage = cfg.MODEL.RESNETS.DEFORM_ON_PER_STAGE641 deform_modulated = cfg.MODEL.RESNETS.DEFORM_MODULATED642 deform_num_groups = cfg.MODEL.RESNETS.DEFORM_NUM_GROUPS643 # fmt: on644 assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)645 646 num_blocks_per_stage = {647 18: [2, 2, 2, 2],648 34: [3, 4, 6, 3],649 50: [3, 4, 6, 3],650 101: [3, 4, 23, 3],651 152: [3, 8, 36, 3],652 }[depth]653 654 if depth in [18, 34]:655 assert out_channels == 64, "Must set MODEL.RESNETS.RES2_OUT_CHANNELS = 64 for R18/R34"656 assert not any(657 deform_on_per_stage658 ), "MODEL.RESNETS.DEFORM_ON_PER_STAGE unsupported for R18/R34"659 assert res5_dilation == 1, "Must set MODEL.RESNETS.RES5_DILATION = 1 for R18/R34"660 assert num_groups == 1, "Must set MODEL.RESNETS.NUM_GROUPS = 1 for R18/R34"661 662 stages = []663 664 for idx, stage_idx in enumerate(range(2, 6)):665 # res5_dilation is used this way as a convention in R-FCN & Deformable Conv paper666 dilation = res5_dilation if stage_idx == 5 else 1667 first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2668 stage_kargs = {669 "num_blocks": num_blocks_per_stage[idx],670 "stride_per_block": [first_stride] + [1] * (num_blocks_per_stage[idx] - 1),671 "in_channels": in_channels,672 "out_channels": out_channels,673 "norm": norm,674 }675 # Use BasicBlock for R18 and R34.676 if depth in [18, 34]:677 stage_kargs["block_class"] = BasicBlock678 else:679 stage_kargs["bottleneck_channels"] = bottleneck_channels680 stage_kargs["stride_in_1x1"] = stride_in_1x1681 stage_kargs["dilation"] = dilation682 stage_kargs["num_groups"] = num_groups683 if deform_on_per_stage[idx]:684 stage_kargs["block_class"] = DeformBottleneckBlock685 stage_kargs["deform_modulated"] = deform_modulated686 stage_kargs["deform_num_groups"] = deform_num_groups687 else:688 stage_kargs["block_class"] = BottleneckBlock689 blocks = ResNet.make_stage(**stage_kargs)690 in_channels = out_channels691 out_channels *= 2692 bottleneck_channels *= 2693 stages.append(blocks)694 return ResNet(stem, stages, out_features=out_features, freeze_at=freeze_at)695