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
