Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import warnings | |
| import torch.nn as nn | |
| from mmengine.model import BaseModule | |
| from mmdet.models.backbones import ResNet | |
| from mmdet.models.layers import ResLayer as _ResLayer | |
| from mmdet.registry import MODELS | |
| class ResLayer(BaseModule): | |
| def __init__(self, | |
| depth, | |
| stage=3, | |
| stride=2, | |
| dilation=1, | |
| style='pytorch', | |
| norm_cfg=dict(type='BN', requires_grad=True), | |
| norm_eval=True, | |
| with_cp=False, | |
| dcn=None, | |
| pretrained=None, | |
| init_cfg=None): | |
| super(ResLayer, self).__init__(init_cfg) | |
| self.norm_eval = norm_eval | |
| self.norm_cfg = norm_cfg | |
| self.stage = stage | |
| self.fp16_enabled = False | |
| block, stage_blocks = ResNet.arch_settings[depth] | |
| stage_block = stage_blocks[stage] | |
| planes = 64 * 2**stage | |
| inplanes = 64 * 2**(stage - 1) * block.expansion | |
| res_layer = _ResLayer( | |
| block, | |
| inplanes, | |
| planes, | |
| stage_block, | |
| stride=stride, | |
| dilation=dilation, | |
| style=style, | |
| with_cp=with_cp, | |
| norm_cfg=self.norm_cfg, | |
| dcn=dcn) | |
| self.add_module(f'layer{stage + 1}', res_layer) | |
| assert not (init_cfg and pretrained), \ | |
| 'init_cfg and pretrained cannot be specified at the same time' | |
| if isinstance(pretrained, str): | |
| warnings.warn('DeprecationWarning: pretrained is a deprecated, ' | |
| 'please use "init_cfg" instead') | |
| self.init_cfg = dict(type='Pretrained', checkpoint=pretrained) | |
| elif pretrained is None: | |
| if init_cfg is None: | |
| self.init_cfg = [ | |
| dict(type='Kaiming', layer='Conv2d'), | |
| dict( | |
| type='Constant', | |
| val=1, | |
| layer=['_BatchNorm', 'GroupNorm']) | |
| ] | |
| else: | |
| raise TypeError('pretrained must be a str or None') | |
| def forward(self, x): | |
| res_layer = getattr(self, f'layer{self.stage + 1}') | |
| out = res_layer(x) | |
| return out | |
| def train(self, mode=True): | |
| super(ResLayer, self).train(mode) | |
| if self.norm_eval: | |
| for m in self.modules(): | |
| if isinstance(m, nn.BatchNorm2d): | |
| m.eval() | |