def __init__( self, backbone: Union[nn.Module, Mapping], neck: Union[nn.Module, Mapping, None] = None, *, pretrained_backbone: Optional[str] = None, output_shapes: List[ShapeSpec], output_names: Optional[List[str]] = None, ): """ Args: backbone: either a backbone module or a mmdet config dict that defines a backbone. The backbone takes a 4D image tensor and returns a sequence of tensors. neck: either a backbone module or a mmdet config dict that defines a neck. The neck takes outputs of backbone and returns a sequence of tensors. If None, no neck is used. pretrained_backbone: defines the backbone weights that can be loaded by mmdet, such as "torchvision://resnet50". output_shapes: shape for every output of the backbone (or neck, if given). stride and channels are often needed. output_names: names for every output of the backbone (or neck, if given). By default, will use "out0", "out1", ... """ super().__init__() if isinstance(backbone, Mapping): from mmdet.models import build_backbone backbone = build_backbone(_to_container(backbone)) self.backbone = backbone if isinstance(neck, Mapping): from mmdet.models import build_neck neck = build_neck(_to_container(neck)) self.neck = neck # It's confusing that backbone weights are given as a separate argument, # but "neck" weights, if any, are part of neck itself. This is the interface # of mmdet so we follow it. Reference: # https://github.com/open-mmlab/mmdetection/blob/master/mmdet/models/detectors/two_stage.py logger.info(f"Initializing mmdet backbone weights: {pretrained_backbone} ...") self.backbone.init_weights(pretrained_backbone) # train() in mmdet modules is non-trivial, and has to be explicitly # called. Reference: # https://github.com/open-mmlab/mmdetection/blob/master/mmdet/models/backbones/resnet.py self.backbone.train() if self.neck is not None: logger.info("Initializing mmdet neck weights ...") if isinstance(self.neck, nn.Sequential): for m in self.neck: m.init_weights() else: self.neck.init_weights() self.neck.train() self._output_shapes = output_shapes if not output_names: output_names = [f"out{i}" for i in range(len(output_shapes))] self._output_names = output_names
def __init__(self, backbone, neck, neck_3d, bbox_head, n_voxels, anchor_generator, train_cfg=None, test_cfg=None, pretrained=None, init_cfg=None): super().__init__(init_cfg=init_cfg) self.backbone = build_backbone(backbone) self.neck = build_neck(neck) self.neck_3d = build_neck(neck_3d) bbox_head.update(train_cfg=train_cfg) bbox_head.update(test_cfg=test_cfg) self.bbox_head = build_head(bbox_head) self.n_voxels = n_voxels self.anchor_generator = build_anchor_generator(anchor_generator) self.train_cfg = train_cfg self.test_cfg = test_cfg
def __init__(self, backbone, neck=None, bbox_head=None, train_cfg=None, test_cfg=None, init_cfg=None, pretrained=None): super(SingleStage3DDetector, self).__init__(init_cfg) self.backbone = build_backbone(backbone) if neck is not None: self.neck = build_neck(neck) bbox_head.update(train_cfg=train_cfg) bbox_head.update(test_cfg=test_cfg) self.bbox_head = build_head(bbox_head) self.train_cfg = train_cfg self.test_cfg = test_cfg
def main(): args = parse_args() config = mmcv.Config.fromfile(os.path.join(root, args.config)) data = torch.randn(1, 3, 800, 800) with torch.no_grad(): backbone = build_backbone(config.model.backbone) neck = build_neck(config.model.neck) rpn_head = build_head(config.model.rpn_head) backbone.eval() neck.eval() rpn_head.eval() #torch.jit.save(rpn_head,'./rpn_head.pt') exit()
def __init__(self, cfg, train_cfg, test_cfg): super(BaseDetector, self).__init__() if 'neck' in cfg: self.neck = build_neck(cfg['neck']) if 'rpn_head' in cfg: self.rpn_head = build_head(cfg['rpn_head']) if 'bbox_head' in cfg: self.bbox_roi_extractor = build_roi_extractor( cfg['bbox_roi_extractor']) self.bbox_head = build_head(cfg['bbox_head']) if 'mask_head' in cfg: self.mask_roi_extractor = build_roi_extractor( cfg['mask_roi_extractor']) self.mask_head = build_head(cfg['mask_head']) self.train_cfg = train_cfg self.test_cfg = test_cfg