def __init__(self, module, wrap_config):
super(CascadeRoIHeadWraper, self).__init__()
self.module = module
self.wrap_config = wrap_config
self.bbox_roi_extractor = [
build_wraper(extractor) for extractor in module.bbox_roi_extractor
]
self.bbox_head = [
build_wraper(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
# init mask if exist
self.enable_mask = False
if 'enable_mask' in wrap_config and wrap_config[
'enable_mask'] and module.with_mask:
self.enable_mask = True
self.init_mask_head(module.mask_roi_extractor, module.mask_head)
self.test_cfg = module.test_cfg
self.num_stages = module.num_stages
def init_mask_head(self, mask_roi_extractor, mask_head):
self.mask_roi_extractor = [
build_wraper(mre) for mre in mask_roi_extractor
]
self.mask_head = [
build_wraper(mh, test_cfg=self.module.test_cfg) for mh in mask_head
]
def __init__(self, module):
super(RPNHeadWraper, self).__init__()
self.module = module
self.anchor_generator = build_wraper(self.module.anchor_generator)
self.bbox_coder = build_wraper(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(TwoStageDetectorWraper, self).__init__()
self.model = model
mmdet_rpn_head = self.model.rpn_head
self.rpn_head_wraper = build_wraper(mmdet_rpn_head, RPNHeadWraper)
mmdet_roi_head = self.model.roi_head
self.roi_head_wraper = build_wraper(mmdet_roi_head,
StandardRoIHeadWraper)
def __init__(self, module):
super(YOLOV3HeadWraper, self).__init__()
self.module = module
self.anchor_generator = build_wraper(self.module.anchor_generator)
self.bbox_coder = build_wraper(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(AnchorHeadWraper, self).__init__()
self.module = module
self.anchor_generator = build_wraper(self.module.anchor_generator)
self.bbox_coder = build_wraper(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(RPNHeadWraper, self).__init__()
self.module = module
self.anchor_generator = build_wraper(self.module.anchor_generator)
self.bbox_coder = build_wraper(self.module.bbox_coder)
self.test_cfg = module.test_cfg
if 'nms' in self.test_cfg:
self.test_cfg.nms_thr = self.test_cfg.nms['iou_threshold']
if 'max_per_img' in self.test_cfg:
self.test_cfg.nms_post = self.test_cfg.max_per_img
self.rpn_nms = BatchedNMS(0.0, self.test_cfg.nms_thr, -1)
def __init__(self, model, wrap_config={}):
super(SingleStageDetectorWraper, self).__init__()
self.model = model
mmdet_backbone = self.model.backbone
self.backbone_wraper = build_wraper(mmdet_backbone, BaseBackboneWraper)
if self.model.with_neck:
mmdet_neck = self.model.neck
self.neck_wraper = build_wraper(mmdet_neck, BaseNeckWraper)
mmdet_bbox_head = self.model.bbox_head
self.bbox_head_wraper = build_wraper(mmdet_bbox_head)
def __init__(self, module):
super(CascadeRoIHeadWraper, self).__init__()
self.module = module
self.bbox_roi_extractor = [build_wraper(extractor) for extractor in module.bbox_roi_extractor]
self.bbox_head = [build_wraper(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(StandardRoIHeadWraper, self).__init__()
self.module = module
self.bbox_roi_extractor = build_wraper(module.bbox_roi_extractor)
self.bbox_head = build_wraper(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(GuidedAnchorHeadWraper, self).__init__()
self.module = module
self.loc_filter_thr = module.loc_filter_thr
self.num_anchors = module.num_anchors
self.square_anchor_generator = build_wraper(self.module.square_anchor_generator)
self.anchor_coder = build_wraper(self.module.anchor_coder)
self.bbox_coder = build_wraper(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(RepPointsHeadWraper, self).__init__(module)
if hasattr(self.module, 'prior_generator'):
# mmdet 2.18
self.prior_generator = build_wraper(self.module.prior_generator)
elif hasattr(self.module, 'point_generators'):
# mmdet 2.10
self.point_generators = [
build_wraper(generator)
for generator in self.module.point_generators
]
else:
self.point_generator = build_wraper(self.module.point_generator)
def __init__(self, module):
super(RepPointsHeadWraper, self).__init__(module)
self.point_generators = [
build_wraper(generator)
for generator in self.module.point_generators
]
def __init__(self, model):
super(TwoStageDetectorWraper, self).__init__()
self.model = model
mmdet_backbone = self.model.backbone
self.backbone_wraper = build_wraper(mmdet_backbone, BaseBackboneWraper)
if self.model.with_neck:
mmdet_neck = self.model.neck
self.neck_wraper = build_wraper(mmdet_neck, BaseNeckWraper)
mmdet_rpn_head = self.model.rpn_head
self.rpn_head_wraper = build_wraper(mmdet_rpn_head, RPNHeadWraper)
mmdet_roi_head = self.model.roi_head
self.roi_head_wraper = build_wraper(mmdet_roi_head,
StandardRoIHeadWraper)
def __init__(self, module):
super(DeformRoiPoolExtractor, self).__init__()
self.module = module
self.roi_layers = [
build_wraper(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(BBoxHeadWraper, self).__init__()
self.module = module
self.bbox_coder = build_wraper(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, wrap_config):
super(HybridTaskCascadeRoIHeadWraper,
self).__init__(module, wrap_config)
module = self.module
self.semantic_head = None
if module.semantic_head is not None:
self.semantic_roi_extractor = build_wraper(
module.semantic_roi_extractor)
self.semantic_head = module.semantic_head
def __init__(self, module, wrap_config={}):
super(StandardRoIHeadWraper, self).__init__()
self.module = module
self.wrap_config = wrap_config
self.bbox_roi_extractor = build_wraper(module.bbox_roi_extractor)
self.bbox_head = build_wraper(module.bbox_head,
test_cfg=module.test_cfg)
if module.with_shared_head:
self.shared_head = module.shared_head
else:
self.shared_head = None
# init mask if exist
self.enable_mask = False
if "enable_mask" in wrap_config and wrap_config[
"enable_mask"] and module.with_mask:
self.enable_mask = True
self.init_mask_head(module.mask_roi_extractor, module.mask_head)
self.test_cfg = module.test_cfg
def __init__(self, module):
super(YOLOXHeadWraper, self).__init__()
self.module = module
self.prior_generator = build_wraper(self.module.prior_generator)
self.cls_out_channels = self.module.cls_out_channels
iou_thr = 0.7
if 'iou_thr' in module.test_cfg.nms:
iou_thr = module.test_cfg.nms.iou_thr
elif 'iou_threshold' in module.test_cfg.nms:
iou_thr = module.test_cfg.nms.iou_threshold
self.test_cfg = module.test_cfg
self.num_classes = self.module.num_classes
self.rcnn_nms = BatchedNMS(module.test_cfg.score_thr,
iou_thr,
backgroundLabelId=-1)
def __init__(self, module):
super(CascadeRPNHeadWraper, self).__init__()
self.module = module
self.stages = [build_wraper(stage) for stage in self.module.stages]
self.test_cfg = module.test_cfg
def __init__(self, model):
super(SingleStageDetectorWraper, self).__init__()
self.model = model
mmdet_bbox_head = self.model.bbox_head
self.bbox_head_wraper = build_wraper(mmdet_bbox_head)
def init_mask_head(self, mask_roi_extractor, mask_head):
self.mask_roi_extractor = build_wraper(mask_roi_extractor)
self.mask_head = build_wraper(mask_head, test_cfg=self.module.test_cfg)
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_wrap_model=False,
output_names=["num_detections", "boxes", "scores", "classes"],
enable_mask=False):
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_wrap_model (bool): return pytorch wrap model, used for debug
output_names (str): the output names of tensorrt engine
enable_mask (bool): weither output the instance segmentation result(w/o postprocess)
"""
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")
if enable_mask and output_names is not None and len(output_names) != 5:
logger.warning("mask mode require len(output_names)==5 " +
"but get output_names=" + str(output_names))
output_names = None
wrap_config = {"enable_mask": enable_mask}
wrap_model = build_wraper(torch_model,
TwoStageDetectorWraper,
wrap_config=wrap_config)
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 = wrap_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_dynamic(
wrap_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_wrap_model:
return trt_model, wrap_model
return trt_model
def __init__(self, module):
super(DETRHeadWraper, self).__init__()
self.module = module
self.test_cfg = module.test_cfg
self.positional_encoding = build_wraper(module.positional_encoding)
def __init__(self, module, wrap_config):
super(GridRoIHeadWraper, self).__init__(module, wrap_config)
self.grid_roi_extractor = build_wraper(module.grid_roi_extractor)
self.grid_head = build_wraper(module.grid_head,
test_cfg=module.test_cfg)