def __init__(self, module): super(RPNHeadWarper, self).__init__() self.module = module self.anchor_generator = build_warper(self.module.anchor_generator) self.bbox_coder = build_warper(self.module.bbox_coder) self.test_cfg = module.test_cfg self.rpn_nms = BatchedNMS(0.0, self.test_cfg.nms_thr, -1)
def __init__(self, model): super(TwoStageDetectorWarper, self).__init__() self.model = model mmdet_rpn_head = self.model.rpn_head self.rpn_head_warper = build_warper(mmdet_rpn_head, RPNHeadWarper) mmdet_roi_head = self.model.roi_head self.roi_head_warper = build_warper(mmdet_roi_head, StandardRoIHeadWarper)
def __init__(self, module): super(AnchorHeadWarper, self).__init__() self.module = module self.anchor_generator = build_warper(self.module.anchor_generator) self.bbox_coder = build_warper(self.module.bbox_coder) self.test_cfg = module.test_cfg self.num_classes = self.module.num_classes self.use_sigmoid_cls = self.module.use_sigmoid_cls self.rcnn_nms = BatchedNMS(module.test_cfg.score_thr, module.test_cfg.nms.iou_threshold, backgroundLabelId = self.num_classes)
def __init__(self, module): super(GuidedAnchorHeadWarper, self).__init__() self.module = module self.loc_filter_thr = module.loc_filter_thr self.num_anchors = module.num_anchors self.square_anchor_generator = build_warper(self.module.square_anchor_generator) self.anchor_coder = build_warper(self.module.anchor_coder) self.bbox_coder = build_warper(self.module.bbox_coder) self.test_cfg = module.test_cfg self.num_classes = self.module.num_classes self.use_sigmoid_cls = self.module.use_sigmoid_cls if "score_thr" in module.test_cfg and "nms" in module.test_cfg and "iou_threshold" in module.test_cfg.nms: self.rcnn_nms = BatchedNMS(module.test_cfg.score_thr, module.test_cfg.nms.iou_threshold, backgroundLabelId = self.num_classes)
def __init__(self, module): super(StandardRoIHeadWarper, self).__init__() self.module = module self.bbox_roi_extractor = build_warper(module.bbox_roi_extractor) self.bbox_head = build_warper(module.bbox_head, test_cfg=module.test_cfg) if module.with_shared_head: self.shared_head = module.shared_head else: self.shared_head = None self.test_cfg = module.test_cfg
def __init__(self, module): super(YOLOV3HeadWarper, self).__init__() self.module = module self.anchor_generator = build_warper(self.module.anchor_generator) self.bbox_coder = build_warper(self.module.bbox_coder) self.featmap_strides = module.featmap_strides self.num_attrib = module.num_attrib self.num_levels = module.num_levels self.test_cfg = module.test_cfg self.num_classes = self.module.num_classes self.rcnn_nms = BatchedNMS(module.test_cfg.score_thr, module.test_cfg.nms.iou_thr, backgroundLabelId=-1)
def __init__(self, module): super(CascadeRoIHeadWarper, self).__init__() self.module = module self.bbox_roi_extractor = [build_warper(extractor) for extractor in module.bbox_roi_extractor] self.bbox_head = [build_warper(bb_head, test_cfg=module.test_cfg) for bb_head in module.bbox_head] if module.with_shared_head: self.shared_head = module.shared_head else: self.shared_head = None self.test_cfg = module.test_cfg self.num_stages = module.num_stages
def __init__(self, module): super(RepPointsHeadWarper, self).__init__(module) self.point_generators = [ build_warper(generator) for generator in self.module.point_generators ]
def __init__(self, module): super(DeformRoiPoolExtractor, self).__init__() self.module = module self.roi_layers = [build_warper(layer) for layer in self.module.roi_layers] self.featmap_strides = self.module.featmap_strides self.finest_scale = self.module.finest_scale
def __init__(self, module, test_cfg): super(BBoxHeadWarper, self).__init__() self.module = module self.bbox_coder = build_warper(self.module.bbox_coder) self.test_cfg = test_cfg self.num_classes = module.num_classes self.rcnn_nms = BatchedNMS(test_cfg.score_thr, test_cfg.nms.iou_threshold, backgroundLabelId = module.num_classes)
def __init__(self, module): super(StandardRoIHeadWarper, self).__init__() self.module = module # self.bbox_roi_extractor = module.bbox_roi_extractor self.bbox_roi_extractor = build_warper(module.bbox_roi_extractor) self.bbox_head = module.bbox_head if module.with_shared_head: self.shared_head = module.shared_head else: self.shared_head = None self.test_cfg = module.test_cfg self.rcnn_nms = BatchedNMS( module.test_cfg.score_thr, module.test_cfg.nms.iou_threshold, backgroundLabelId=self.bbox_head.num_classes)
def __init__(self, module): super(CascadeRoIHeadWarper, self).__init__() self.module = module self.bbox_roi_extractor = [ build_warper(extractor) for extractor in module.bbox_roi_extractor ] self.bbox_head = module.bbox_head if module.with_shared_head: self.shared_head = module.shared_head else: self.shared_head = None self.test_cfg = module.test_cfg self.num_stages = module.num_stages self.rcnn_nms = BatchedNMS( module.test_cfg.score_thr, module.test_cfg.nms.iou_threshold, backgroundLabelId=self.bbox_head[-1].num_classes)
def __init__(self, model): super(SingleStageDetectorWarper, self).__init__() self.model = model mmdet_bbox_head = self.model.bbox_head self.bbox_head_warper = build_warper(mmdet_bbox_head)
def mmdet2trt( config, checkpoint, device="cuda:0", fp16_mode=False, max_workspace_size=0.5e9, opt_shape_param=None, trt_log_level="INFO", return_warp_model=False, output_names=["num_detections", "boxes", "scores", "classes"], ): device = torch.device(device) logger.info("Loading model from config: {}".format(config)) torch_model = init_detector(config, checkpoint=checkpoint, device=device) cfg = torch_model.cfg logger.info("Wrapping model") warp_model = build_warper(torch_model, TwoStageDetectorWarper) if opt_shape_param is None: img_scale = cfg.test_pipeline[1]["img_scale"] min_scale = min(img_scale) max_scale = max(img_scale)+32 opt_shape_param = [ [ [1, 3, min_scale, min_scale], [1, 3, img_scale[1], img_scale[0]], [1, 3, max_scale, max_scale], ] ] dummy_shape = opt_shape_param[0][1] dummy_input = torch.rand(dummy_shape).to(device) logger.info("Model warmup") with torch.no_grad(): result = warp_model(dummy_input) logger.info("Converting model") start = time.time() with torch.cuda.device(device), torch.no_grad(): trt_model = torch2trt( warp_model, [dummy_input], log_level=getattr(trt.Logger, trt_log_level), fp16_mode=fp16_mode, opt_shape_param=opt_shape_param, max_workspace_size=max_workspace_size, keep_network=False, strict_type_constraints=True, output_names=output_names, ) duration = time.time() - start logger.info("Conversion took {} s".format(duration)) if return_warp_model: return trt_model, warp_model return trt_model
def mmdet2trt( config, checkpoint, device="cuda:0", fp16_mode=False, int8_mode=False, int8_calib_dataset=None, int8_calib_alg="entropy", max_workspace_size=0.5e9, opt_shape_param=None, trt_log_level="INFO", return_warp_model=False, output_names=["num_detections", "boxes", "scores", "classes"], ): r""" create tensorrt model from mmdetection. Args: config (str): config file path of mmdetection model checkpoint (str): checkpoint file path of mmdetection model device (str): convert gpu device fp16_mode (bool): create fp16 mode engine. int8_mode (bool): create int8 mode engine. int8_calib_dataset (object): dataset object used to do data calibrate int8_calib_alg (str): how to calibrate int8, ["minmax", "entropy"] max_workspace_size (int): tensorrt workspace size. some tactic might need large workspace. opt_shape_param (list[list[list[int]]]): the min/optimize/max shape of input tensor trt_log_level (str): tensorrt log level, ["VERBOSE", "INFO", "WARNING", "ERROR"] return_warp_model (bool): return pytorch warp model, used for debug output_names (str): the output names of tensorrt engine """ device = torch.device(device) logger.info("Loading model from config: {}".format(config)) torch_model = init_detector(config, checkpoint=checkpoint, device=device) cfg = torch_model.cfg logger.info("Wrapping model") warp_model = build_warper(torch_model, TwoStageDetectorWarper) if opt_shape_param is None: img_scale = cfg.test_pipeline[1]["img_scale"] min_scale = min(img_scale) max_scale = max(img_scale)+32 opt_shape_param = [ [ [1, 3, min_scale, min_scale], [1, 3, img_scale[1], img_scale[0]], [1, 3, max_scale, max_scale], ] ] dummy_shape = opt_shape_param[0][1] dummy_input = torch.rand(dummy_shape).to(device) dummy_input = (dummy_input-0.45)/0.27 dummy_input = dummy_input.contiguous() logger.info("Model warmup") with torch.no_grad(): result = warp_model(dummy_input) logger.info("Converting model") start = time.time() with torch.cuda.device(device), torch.no_grad(): int8_calib_algorithm = trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2 if int8_calib_alg=="minmax": int8_calib_algorithm = trt.CalibrationAlgoType.MINMAX_CALIBRATION elif int8_calib_alg=="entropy": int8_calib_algorithm = trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2 trt_model = torch2trt( warp_model, [dummy_input], log_level=getattr(trt.Logger, trt_log_level), fp16_mode=fp16_mode, opt_shape_param=opt_shape_param, max_workspace_size=int(max_workspace_size), keep_network=False, strict_type_constraints=True, output_names=output_names, int8_mode=int8_mode, int8_calib_dataset=int8_calib_dataset, int8_calib_algorithm=int8_calib_algorithm ) duration = time.time() - start logger.info("Conversion took {} s".format(duration)) if return_warp_model: return trt_model, warp_model return trt_model