def convert_AnchorGeneratorDynamic(ctx):
module = ctx.method_args[0]
input = ctx.method_args[1]
input_trt = trt_(ctx.network, input)
output = ctx.method_return
ag = module
index = ag.index
base_size = ag.base_size
stride = get_arg(ctx, 'stride', pos=2, default=base_size)
if hasattr(ag.generator, 'base_anchors'):
base_anchors = ag.generator.base_anchors[index]
# base_anchors = base_anchors.view(-1).cpu().numpy()
base_anchors_trt = trt_(ctx.network, base_anchors.float())
plugin = create_gridanchordynamic_plugin('ag_' + str(id(module)),
stride=stride)
else:
print('no base_anchors in {}'.format(ag.generator))
# scales = ag.scales.detach().cpu().numpy().astype(np.float32)
# ratios = ag.ratios.detach().cpu().numpy().astype(np.float32)
# scale_major = ag.scale_major
# ctr = ag.ctr
# if ctr is None:
# # center_x = -1
# # center_y = -1
# center_x = 0
# center_y = 0
# else:
# center_x, center_y = ag.ctr
# plugin = create_gridanchordynamic_plugin("ag_" + str(id(module)),
# base_size=base_size,
# stride=stride,
# scales=scales,
# ratios=ratios,
# scale_major=scale_major,
# center_x=center_x,
# center_y=center_y)
custom_layer = ctx.network.add_plugin_v2(
inputs=[input_trt, base_anchors_trt], plugin=plugin)
output._trt = custom_layer.get_output(0)
def convert_batchednms(ctx):
module = ctx.method_args[0]
scores = ctx.method_args[1]
bboxes = ctx.method_args[2]
topK = ctx.method_args[3]
keepTopK = ctx.method_args[4]
topK = min(scores.shape[1], topK)
keepTopK = min(scores.shape[1], keepTopK)
scoreThreshold = module.scoreThreshold
iouThreshold = module.iouThreshold
backgroundLabelId = module.backgroundLabelId
numClasses = scores.shape[2]
shareLocation = (bboxes.shape[2] == 1)
scores_trt = trt_(ctx.network, scores)
bboxes_trt = trt_(ctx.network, bboxes)
output = ctx.method_return
plugin = create_batchednms_plugin("batchednms_" + str(id(module)),
scoreThreshold,
iouThreshold,
topK,
keepTopK,
numClasses,
backgroundLabelId,
shareLocation=shareLocation,
isNormalized=False,
clipBoxes=False)
custom_layer = ctx.network.add_plugin_v2(inputs=[bboxes_trt, scores_trt],
plugin=plugin)
if isinstance(output, torch.Tensor):
output._trt = custom_layer.get_output(0)
else:
for i in range(len(output)):
output[i]._trt = custom_layer.get_output(i)
def convert_delta2bbox(ctx):
cls_scores = get_arg(ctx, 'cls_scores', pos=0, default=None)
bbox_preds = get_arg(ctx, 'bbox_preds', pos=1, default=None)
anchors = get_arg(ctx, 'anchors', pos=2, default=None)
min_num_bboxes = get_arg(ctx, 'min_num_bboxes', pos=3, default=1000)
num_classes = get_arg(ctx, 'num_classes', pos=4, default=1)
use_sigmoid_cls = get_arg(ctx, 'use_sigmoid_cls', pos=5, default=True)
target_mean = get_arg(ctx, 'target_mean', pos=6, default=[0, 0, 0, 0])
target_std = get_arg(ctx, 'target_std', pos=7, default=[1, 1, 1, 1])
input_x = get_arg(ctx, 'input_x', pos=8, default=None)
scores_trt = trt_(ctx.network, cls_scores)
preds_trt = trt_(ctx.network, bbox_preds)
anchors_trt = trt_(ctx.network, anchors)
if input_x is not None:
input_x_trt = trt_(ctx.network, input_x)
input_x_shape_trt = ctx.network.add_shape(input_x_trt).get_output(0)
output = ctx.method_return
plugin = create_delta2bbox_custom_plugin("delta2bbox_custom_" +
str(id(cls_scores)),
use_sigmoid_cls=use_sigmoid_cls,
min_num_bbox=min_num_bboxes,
num_classes=num_classes,
target_means=target_mean,
target_stds=target_std)
layer_input = [scores_trt, preds_trt, anchors_trt]
if input_x is not None:
layer_input.append(input_x_shape_trt)
custom_layer = ctx.network.add_plugin_v2(inputs=layer_input, plugin=plugin)
if isinstance(output, torch.Tensor):
output._trt = custom_layer.get_output(0)
else:
for i in range(len(output)):
output[i]._trt = custom_layer.get_output(i)
def convert_AnchorGeneratorDynamic(ctx):
module = ctx.method_args[0]
input = ctx.method_args[1]
input_trt = trt_(ctx.network, input)
output = ctx.method_return
ag = module
index = ag.index
base_size = ag.base_size
stride = get_arg(ctx, 'stride', pos=2, default=base_size)
if hasattr(ag.generator, 'base_anchors'):
base_anchors = ag.generator.base_anchors[index]
base_anchors = base_anchors.view(-1).cpu().numpy()
plugin = create_gridanchordynamic_plugin("ag_" + str(id(module)),
base_size=base_size,
stride=stride,
base_anchors=base_anchors)
else:
scales = ag.scales.detach().cpu().numpy().astype(np.float32)
ratios = ag.ratios.detach().cpu().numpy().astype(np.float32)
scale_major = ag.scale_major
ctr = ag.ctr
if ctr is None:
# center_x = -1
# center_y = -1
center_x = 0
center_y = 0
else:
center_x, center_y = ag.ctr
plugin = create_gridanchordynamic_plugin("ag_" + str(id(module)),
base_size=base_size,
stride=stride,
scales=scales,
ratios=ratios,
scale_major=scale_major,
center_x=center_x,
center_y=center_y)
custom_layer = ctx.network.add_plugin_v2(inputs=[input_trt], plugin=plugin)
output._trt = custom_layer.get_output(0)