def convert_prelu(ctx):
input = get_arg(ctx, 'input', pos=0, default=None)
weight = get_arg(ctx, 'weight', pos=1, default=None)
output = ctx.method_return
weight_shape = [1] * len(input.shape)
weight_shape[1] = weight.numel()
input_trt = trt_(ctx.network, input)
# y = prelu(x) = relu(x) - alpha * relu(-x)
weight_trt = ctx.network.add_constant(
weight_shape,
-weight.detach().view(weight_shape).cpu().numpy()).get_output(
0) # detach so considered leaf
# x >= 0
a = ctx.network.add_activation(input_trt,
trt.ActivationType.RELU).get_output(0)
# x <= 0
b = ctx.network.add_unary(input_trt, trt.UnaryOperation.NEG).get_output(0)
b = ctx.network.add_activation(b, trt.ActivationType.RELU).get_output(0)
b = ctx.network.add_elementwise(
b, weight_trt, trt.ElementWiseOperation.PROD).get_output(0)
# y = a + b
y = ctx.network.add_elementwise(a, b, trt.ElementWiseOperation.SUM)
output._trt = y.get_output(0)
def convert_carafe_tensor_add(ctx):
a = get_arg(ctx, 'a', pos=1, default=None)
b = get_arg(ctx, 'b', pos=2, default=None)
a_trt = trt_(ctx.network, a)
b_trt = trt_(ctx.network, b)
output = ctx.method_return
a_shape_trt = ctx.network.add_shape(a_trt).get_output(0)
layer = ctx.network.add_slice(b_trt, [0] * len(a.shape),
[2] * len(a.shape), [1] * len(a.shape))
layer.set_input(
1,
trt_(ctx.network,
torch.tensor([0] * len(a.shape), dtype=torch.int32).to(a.device)))
layer.set_input(2, a_shape_trt)
layer.set_input(
3,
trt_(ctx.network,
torch.tensor([1] * len(a.shape), dtype=torch.int32).to(a.device)))
new_b_trt = layer.get_output(0)
layer = ctx.network.add_elementwise(a_trt, new_b_trt,
trt.ElementWiseOperation.SUM)
output._trt = layer.get_output(0)
def convert_roiextractor(ctx):
module = ctx.method_args[0]
feats = get_arg(ctx, 'feats', pos=1, default=None)
rois = get_arg(ctx, 'rois', pos=2, default=None)
roi_scale_factor = get_arg(ctx, 'roi_scale_factor', pos=3, default=None)
if not roi_scale_factor:
roi_scale_factor = -1.0
out_size = module.roi_layers[0].output_size[0]
sample_num = module.roi_layers[0].sampling_ratio
featmap_strides = module.featmap_strides
finest_scale = module.finest_scale
feats_trt = [trt_(ctx.network, f) for f in feats]
rois_trt = trt_(ctx.network, rois)
output = ctx.method_return
plugin = create_roiextractor_plugin(
'roiextractor_' + str(id(module)),
out_size,
sample_num,
featmap_strides,
roi_scale_factor,
finest_scale,
aligned=1)
custom_layer = ctx.network.add_plugin_v2(
inputs=[rois_trt] + feats_trt, plugin=plugin)
output._trt = custom_layer.get_output(0)
def convert_mean(ctx):
input = ctx.method_args[0]
input_trt = trt_(ctx.network, input)
output = ctx.method_return
dim = get_arg(ctx, 'dim', pos=1, default=None)
keep_dims = get_arg(ctx, 'keepdim', pos=2, default=False)
# get dims from args or kwargs
if dim is None:
dim = tuple(range(len(input.shape)))
# convert list to tuple
if isinstance(dim, list):
dim = tuple(dim)
if not isinstance(dim, tuple):
dim = (dim, )
dim = tuple([d if d >= 0 else len(input.shape) + d for d in dim])
# create axes bitmask for reduce layer
axes = 0
for d in dim:
axes |= 1 << d
layer = ctx.network.add_reduce(input_trt, trt.ReduceOperation.AVG, axes,
keep_dims)
output._trt = layer.get_output(0)
def convert_split(ctx):
input = get_arg(ctx, 'input', 0, None)
input_trt = trt_(ctx.network, input)
# we don't need to parse split/chunk (arg 1)
# since we infer size from output tensors
dim = get_arg(ctx, 'dim', 2, 0)
outputs = ctx.method_return
assert (dim >= 1)
start = [0] * len(input.shape) # exclude batch
stride = [1] * len(start)
offset = 0
trt_dim = dim
# add slice layers
for i, output in enumerate(outputs):
shape = list(output.shape)
start[trt_dim] = offset
layer = ctx.network.add_slice(input_trt,
start=start,
shape=shape,
stride=stride)
output._trt = layer.get_output(0)
offset = offset + shape[trt_dim]
def convert_Tensor_where(ctx):
x = ctx.method_args[0]
condition = get_arg(ctx, 'condition', pos=1, default=None)
y = get_arg(ctx, 'y', pos=2, default=None)
ctx.method_args = [condition, x, y]
ctx.method_kwargs = {}
convert_where(ctx)
def convert_prod(ctx):
input = ctx.method_args[0]
dim = get_arg(ctx, 'dim', pos=1, default=tuple(range(1, input.ndim)))
keepdim = get_arg(ctx, 'keepdim', pos=2, default=False)
input_trt = trt_(ctx.network, input)
output = ctx.method_return
layer = ctx.network.add_reduce(input_trt, trt.ReduceOperation.PROD,
torch_dim_to_trt_axes(dim), keepdim)
output._trt = layer.get_output(0)
def convert_elu(ctx):
input = get_arg(ctx, 'input', pos=0, default=None)
alpha = get_arg(ctx, 'alpha', pos=1, default=1.0)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
layer = ctx.network.add_activation(input_trt, trt.ActivationType.ELU)
layer.alpha = alpha
output._trt = layer.get_output(0)
def convert_leaky_relu(ctx):
input = get_arg(ctx, 'input', pos=0, default=None)
negative_slope = get_arg(ctx, 'negative_slope', pos=1, default=0.01)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
layer = ctx.network.add_activation(input_trt,
trt.ActivationType.LEAKY_RELU)
layer.alpha = negative_slope
output._trt = layer.get_output(0)
def convert_max_pool1d(ctx):
# parse args
old_args = ctx.method_args
old_kwargs = ctx.method_kwargs
input = get_arg(ctx, 'input', pos=0, default=None)
kernel_size = get_arg(ctx, 'kernel_size', pos=1, default=None)
stride = get_arg(ctx, 'stride', pos=2, default=None)
padding = get_arg(ctx, 'padding', pos=3, default=0)
# dilation = get_arg(ctx, 'dilation', pos=4, default=1)
ceil_mode = get_arg(ctx, 'ceil_mode', pos=5, default=False)
kernel_size = (kernel_size, 1)
stride = (stride, 1)
padding = (padding, 0)
output = ctx.method_return
# unsqueeze -1
unsqueeze_input = input.unsqueeze(-1)
ctx.method_args = [input, -1]
ctx.method_kwargs = {}
ctx.method_return = unsqueeze_input
convert_unsqueeze(ctx)
# pool2d
input_trt = trt_(ctx.network, unsqueeze_input)
layer = ctx.network.add_pooling(input=input_trt,
type=trt.PoolingType.MAX,
window_size=kernel_size)
layer.stride = stride
layer.padding = padding
if ceil_mode:
layer.padding_mode = trt.PaddingMode.EXPLICIT_ROUND_UP
pool2d_output = torch.nn.functional.max_pool2d(unsqueeze_input,
kernel_size=kernel_size,
stride=stride,
padding=padding,
ceil_mode=ceil_mode)
pool2d_output._trt = layer.get_output(0)
# squeeze -1
ctx.method_args = [pool2d_output, -1]
ctx.method_kwargs = {}
ctx.method_return = output
convert_squeeze(ctx)
ctx.method_args = old_args
ctx.method_kwargs = old_kwargs
ctx.method_return = output
def convert_where(ctx):
condition = get_arg(ctx, 'condition', pos=0, default=None)
x = get_arg(ctx, 'x', pos=1, default=None)
y = get_arg(ctx, 'y', pos=2, default=None)
condition_trt = trt_(ctx.network, condition)
x_trt = trt_(ctx.network, x)
y_trt = trt_(ctx.network, y)
output = ctx.method_return
layer = ctx.network.add_select(condition_trt, x_trt, y_trt)
output_trt = layer.get_output(0)
output._trt = output_trt
def convert_narrow(ctx):
input = ctx.method_args[0]
if 'dim' in ctx.method_kwargs:
dim = ctx.method_kwargs['dim']
elif 'dimension' in ctx.method_kwargs:
dim = ctx.method_kwargs['dimension']
else:
dim = ctx.method_args[1]
input_dim = input.dim()
if dim < 0:
dim = dim + input_dim
start = get_arg(ctx, 'start', pos=2, default=None)
length = get_arg(ctx, 'length', pos=3, default=None)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
input_shape_trt = tensor_trt_get_shape_trt(ctx.network, input_trt)
start_trt = get_intwarper_trt(start, ctx)
length_trt = get_intwarper_trt(length, ctx)
stride_trt = trt_(ctx.network, torch.ones([input_dim]).int())
if dim != 0:
start_pre_trt = trt_(ctx.network,
torch.zeros([
dim,
]).int())
start_trt = ctx.network.add_concatenation([start_pre_trt,
start_trt]).get_output(0)
length_pre_trt = slice_shape_trt(ctx.network, input_shape_trt, 0, dim)
length_trt = ctx.network.add_concatenation(
[length_pre_trt, length_trt]).get_output(0)
if dim < input_dim - 1:
start_post_trt = trt_(ctx.network,
torch.zeros([input_dim - dim - 1]).int())
start_trt = ctx.network.add_concatenation([start_trt, start_post_trt
]).get_output(0)
length_post_trt = slice_shape_trt(ctx.network, input_shape_trt,
dim + 1)
length_trt = ctx.network.add_concatenation(
[length_trt, length_post_trt]).get_output(0)
layer = ctx.network.add_slice(input_trt, [0] * input_dim, [1] * input_dim,
[1] * input_dim)
layer.set_input(1, start_trt)
layer.set_input(2, length_trt)
layer.set_input(3, stride_trt)
output._trt = layer.get_output(0)
def convert_adaptive_avg_pool2d(ctx):
input = ctx.method_args[0]
output_size = get_arg(ctx, 'output_size', pos=1, default=0)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
if isinstance(output_size, int):
output_size = (output_size, output_size)
output_size = tuple([-1 if not o else o for o in output_size])
if output_size[0] == 1 and output_size[1] == 1:
# use reduce as max pool2d
shape_length = len(input.shape)
axes = (1 << (shape_length - 1)) + (1 << (shape_length - 2))
keepdim = True
layer = ctx.network.add_reduce(input_trt, trt.ReduceOperation.AVG,
axes, keepdim)
output._trt = layer.get_output(0)
else:
plugin = create_adaptivepool_plugin(
'adaptive_avg_pool2d_' + str(id(input)),
output_size=output_size,
pooling_type=trt.PoolingType.AVERAGE)
layer = ctx.network.add_plugin_v2(inputs=[input_trt], plugin=plugin)
output._trt = layer.get_output(0)
def convert_clamp(ctx):
input = ctx.method_args[0]
min_val = get_arg(ctx, 'min', pos=1, default=None)
max_val = get_arg(ctx, 'max', pos=2, default=None)
input_trt = trt_(ctx.network, input)
output = ctx.method_return
if min_val is not None:
layer = __add_clamp(ctx.network, input_trt, min_val,
trt.ElementWiseOperation.MAX)
input_trt = layer.get_output(0)
if max_val is not None:
layer = __add_clamp(ctx.network, input_trt, max_val,
trt.ElementWiseOperation.MIN)
output._trt = layer.get_output(0)
def convert_squeeze(ctx):
input = ctx.method_args[0]
dim = get_arg(ctx, 'dim', pos=1, default=None)
if dim is None:
dim = list(
filter(lambda x: input.shape[x] == 1, range(len(input.shape))))
else:
if input.shape[dim] != 1:
ctx.method_args = [input]
convert_identity(ctx)
return
if dim < 0:
dim = len(input.shape) + dim
dim = [dim]
input_trt = trt_(ctx.network, input)
shape_trt = ctx.network.add_shape(input_trt).get_output(0)
output = ctx.method_return
reverse_dim = list(filter(lambda x: x not in dim, range(len(input.shape))))
reverse_dim_trt = trt_(
ctx.network,
torch.tensor(reverse_dim, dtype=torch.int32).to(input.device))
new_shape_trt = ctx.network.add_gather(shape_trt, reverse_dim_trt,
0).get_output(0)
layer = ctx.network.add_shuffle(input_trt)
layer.set_input(1, new_shape_trt)
output._trt = layer.get_output(0)
def convert_topk(ctx):
input = ctx.method_args[0]
k = get_arg(ctx, 'k', pos=1, default=1)
axis = get_arg(ctx, 'dim', pos=2, default=len(input.shape) - 1)
if axis is None:
axis = len(input.shape) - 1
if axis < 0:
axis = len(input.shape) + axis
if k > 3840:
print('warning: topk = ' + k +
' > 3840 is not allowed in TensorRT, use 3840 instead.')
k = 3840
largest = get_arg(ctx, 'largest', pos=3, default=True)
topkOp = trt.TopKOperation.MAX if largest else trt.TopKOperation.MIN
input_trt = trt_(ctx.network, input)
output = ctx.method_return
# can only use topk on dim>=2
need_unsqueeze = len(input_trt.shape) == 1
if need_unsqueeze:
layer = ctx.network.add_shuffle(input_trt)
layer.reshape_dims = (1, ) + tuple(input_trt.shape)
input_trt = layer.get_output(0)
axis += 1
layer = ctx.network.add_topk(input_trt, topkOp, k, 1 << axis)
output0_trt = layer.get_output(0)
output1_trt = layer.get_output(1)
# recovery
if need_unsqueeze:
layer = ctx.network.add_shuffle(output0_trt)
layer.reshape_dims = tuple(output0_trt.shape)[1:]
output0_trt = layer.get_output(0)
layer = ctx.network.add_shuffle(output1_trt)
layer.reshape_dims = tuple(output1_trt.shape)[1:]
output1_trt = layer.get_output(0)
output[0]._trt = output0_trt
output[1]._trt = output1_trt
def convert_RoIPool(ctx):
module = ctx.method_args[0]
input = get_arg(ctx, 'input', pos=1, default=None)
boxes = get_arg(ctx, 'boxes', pos=2, default=None)
output_size = module.output_size
spatial_scale = module.spatial_scale
old_method_args = ctx.method_args
old_method_kwargs = ctx.method_kwargs
new_method_args = [input, boxes, output_size, spatial_scale]
new_method_kwargs = {}
ctx.method_args = new_method_args
ctx.method_kwargs = new_method_kwargs
convert_roi_pool(ctx)
ctx.method_args = old_method_args
ctx.method_kwargs = old_method_kwargs
def convert_softsign(ctx):
input = get_arg(ctx, 'input', pos=0, default=None)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
layer = ctx.network.add_activation(input_trt, trt.ActivationType.SOFTSIGN)
output._trt = layer.get_output(0)
def convert_ModulatedDeformConv(ctx):
input = get_arg(ctx, 'input', pos=0, default=None)
offset = get_arg(ctx, 'offset', pos=1, default=None)
mask = get_arg(ctx, 'mask', pos=2, default=None)
weight = get_arg(ctx, 'weight', pos=3, default=None)
bias = get_arg(ctx, 'bias', pos=4, default=None)
stride = get_arg(ctx, 'stride', pos=5, default=1)
padding = get_arg(ctx, 'padding', pos=6, default=0)
dilation = get_arg(ctx, 'dilation', pos=7, default=1)
groups = get_arg(ctx, 'groups', pos=8, default=1)
deform_groups = get_arg(ctx, 'deform_groups', pos=9, default=1)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
offset_trt = trt_(ctx.network, offset)
mask_trt = trt_(ctx.network, mask)
weight_trt = trt_(ctx.network, weight)
if bias is not None:
bias_trt = trt_(ctx.network, bias)
if not isinstance(stride, tuple):
stride = (stride, ) * 2
if not isinstance(padding, tuple):
padding = (padding, ) * 2
if not isinstance(dilation, tuple):
dilation = (dilation, ) * 2
plugin = create_dcnv2_plugin('dcn_' + str(id(input)),
stride=stride,
padding=padding,
dilation=dilation,
deformable_group=deform_groups,
group=groups)
layer_inputs = [input_trt, offset_trt, mask_trt, weight_trt]
if bias is not None:
layer_inputs += [bias_trt]
custom_layer = ctx.network.add_plugin_v2(inputs=layer_inputs,
plugin=plugin)
output._trt = custom_layer.get_output(0)
def convert_DeformPool(ctx):
input = get_arg(ctx, 'input', pos=0, default=None)
rois = get_arg(ctx, 'rois', pos=1, default=None)
offset = get_arg(ctx, 'offset', pos=2, default=None)
out_size = get_arg(ctx, 'output_size', pos=3, default=(7, 7))
spatial_scale = get_arg(ctx, 'spatial_scale', pos=4, default=1.)
sampling_ratio = get_arg(ctx, 'sampling_ratio', pos=5, default=0)
gamma = get_arg(ctx, 'gamma', pos=6, default=0.1)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
rois_trt = trt_(ctx.network, rois)
offset_trt = None
if offset is not None and len(offset.shape) > 1:
offset_trt = trt_(ctx.network, offset)
plugin = create_deformable_pool_plugin('deform_roi_pool_' + str(id(input)),
out_size, spatial_scale,
sampling_ratio, gamma)
if offset_trt is None:
custom_layer = ctx.network.add_plugin_v2(
inputs=[input_trt, rois_trt], plugin=plugin)
else:
custom_layer = ctx.network.add_plugin_v2(
inputs=[input_trt, rois_trt, offset_trt], plugin=plugin)
output._trt = custom_layer.get_output(0)
def convert_carafe_kernel_normalizer(ctx):
import torch
from torch.nn import functional as F
module = ctx.method_args[0]
mask = get_arg(ctx, 'mask', pos=1, default=None)
scale_factor = module.scale_factor
up_kernel = module.up_kernel
output = ctx.method_return
# pixel shuffle
ps_mask = F.pixel_shuffle(mask, scale_factor)
ctx.method_args = [mask, scale_factor]
ctx.method_return = ps_mask
convert_pixel_shuffle(ctx)
# view0
n, mask_c, h, w = ps_mask.size()
mask_channel = int(mask_c / (up_kernel * up_kernel))
view_ps_mask = ps_mask.view(n, mask_channel, -1, h, w)
ps_mask_trt = trt_(ctx.network, ps_mask)
ps_mask_shape_trt = ctx.network.add_shape(ps_mask_trt).get_output(0)
ps_mask_batch_trt = ctx.network.add_slice(ps_mask_shape_trt, [0], [1],
[1]).get_output(0)
ps_mask_channel_trt = ctx.network.add_slice(ps_mask_shape_trt, [1], [1],
[1]).get_output(0)
ps_mask_hw_trt = ctx.network.add_slice(ps_mask_shape_trt, [2], [2],
[1]).get_output(0)
kernel_v2_trt = trt_(
ctx.network,
torch.tensor([up_kernel * up_kernel],
dtype=torch.int32).to(mask.device))
ps_mask_new_channel_trt = ctx.network.add_elementwise(
ps_mask_channel_trt, kernel_v2_trt,
trt.ElementWiseOperation.FLOOR_DIV).get_output(0)
ps_mask_new_shape_trt = ctx.network.add_concatenation([
ps_mask_batch_trt, ps_mask_new_channel_trt, kernel_v2_trt,
ps_mask_hw_trt
]).get_output(0)
layer = ctx.network.add_shuffle(ps_mask_trt)
layer.set_input(1, ps_mask_new_shape_trt)
view_ps_mask._trt = layer.get_output(0)
# softmax
softmax_mask = F.softmax(view_ps_mask, dim=2)
ctx.method_args = [view_ps_mask, 2]
ctx.method_return = softmax_mask
convert_softmax(ctx)
# view1
softmax_mask_trt = trt_(ctx.network, softmax_mask)
layer = ctx.network.add_shuffle(softmax_mask_trt)
layer.set_input(1, ps_mask_shape_trt)
output._trt = layer.get_output(0)
ctx.method_return = output
def convert_adaptive_max_pool2d_by_input(ctx):
input = get_arg(ctx, 'x', pos=0, default=None)
shape_wraper = get_arg(ctx, 'shape_wraper', pos=1, default=None)
output = ctx.method_return
output_size = shape_wraper.shape
input_trt = trt_(ctx.network, input)
wrapshape_trt = trt_(ctx.network, shape_wraper)
plugin = create_adaptivepool_plugin('adaptive_max_pool2d_by_input_' +
str(id(input)),
output_size=output_size,
pooling_type=trt.PoolingType.MAX)
layer = ctx.network.add_plugin_v2(inputs=[input_trt, wrapshape_trt],
plugin=plugin)
output._trt = layer.get_output(0)
def convert_nms(ctx):
boxes = get_arg(ctx, 'boxes', pos=0, default=None)
scores = get_arg(ctx, 'scores', pos=1, default=None)
iou_threshold = get_arg(ctx, 'iou_threshold', pos=2, default=0.7)
output = ctx.method_return
boxes_trt = trt_(ctx.network, boxes)
scores_trt = trt_(ctx.network, scores)
plugin = create_nms_plugin('nms_' + str(id(boxes)),
iou_threshold=iou_threshold)
custom_layer = ctx.network.add_plugin_v2(inputs=[boxes_trt, scores_trt],
plugin=plugin)
output._trt = custom_layer.get_output(0)
def convert_unfold(ctx):
input = ctx.method_args[0]
kernel_size = get_arg(ctx, 'kernel_size', pos=1, default=0)
dilation = get_arg(ctx, 'dilation', pos=2, default=1)
padding = get_arg(ctx, 'padding', pos=3, default=0)
stride = get_arg(ctx, 'stride', pos=4, default=1)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
plugin = create_torchunfold_plugin('unfold_' + str(id(input)),
kernel_size=kernel_size,
dilation=dilation,
padding=padding,
stride=stride)
layer = ctx.network.add_plugin_v2(inputs=[input_trt], plugin=plugin)
output._trt = layer.get_output(0)
def convert_clamp_max(ctx):
input = ctx.method_args[0]
input_trt = trt_(ctx.network, input)
val = get_arg(ctx, 'max', pos=1, default=0)
output = ctx.method_return
layer = __add_clamp(ctx.network, input_trt, val,
trt.ElementWiseOperation.MIN)
output._trt = layer.get_output(0)
def convert_pixel_shuffle(ctx):
input = ctx.method_args[0]
upscale_factor = get_arg(ctx, 'upscale_factor', pos=1, default=None)
input_trt = trt_(ctx.network, input)
input_shape_trt = ctx.network.add_shape(input_trt).get_output(0)
output = ctx.method_return
batch_shape_trt = ctx.network.add_slice(input_shape_trt, [0], [1],
[1]).get_output(0)
channel_shape_trt = ctx.network.add_slice(input_shape_trt, [1], [1],
[1]).get_output(0)
height_shape_trt = ctx.network.add_slice(input_shape_trt, [2], [1],
[1]).get_output(0)
width_shape_trt = ctx.network.add_slice(input_shape_trt, [3], [1],
[1]).get_output(0)
upscale_shape_trt = trt_(
ctx.network,
torch.tensor([upscale_factor], dtype=torch.int32).to(input.device))
upscale_p2_trt = ctx.network.add_elementwise(
upscale_shape_trt, upscale_shape_trt,
trt.ElementWiseOperation.PROD).get_output(0)
new_channel_shape_trt = ctx.network.add_elementwise(
channel_shape_trt, upscale_p2_trt,
trt.ElementWiseOperation.FLOOR_DIV).get_output(0)
# (b, c0, s, s, h, w)
pre_shape_trt = ctx.network.add_concatenation([
batch_shape_trt, new_channel_shape_trt, upscale_shape_trt,
upscale_shape_trt, height_shape_trt, width_shape_trt
]).get_output(0)
layer = ctx.network.add_shuffle(input_trt)
layer.set_input(1, pre_shape_trt)
layer.second_transpose = (0, 1, 4, 2, 5, 3)
permute_trt = layer.get_output(0)
new_height_shape_trt = ctx.network.add_elementwise(
height_shape_trt, upscale_shape_trt,
trt.ElementWiseOperation.PROD).get_output(0)
new_width_shape_trt = ctx.network.add_elementwise(
width_shape_trt, upscale_shape_trt,
trt.ElementWiseOperation.PROD).get_output(0)
post_shape_trt = ctx.network.add_concatenation([
batch_shape_trt, new_channel_shape_trt, new_height_shape_trt,
new_width_shape_trt
]).get_output(0)
layer = ctx.network.add_shuffle(permute_trt)
layer.set_input(1, post_shape_trt)
output._trt = layer.get_output(0)
def convert_selu(ctx):
input = get_arg(ctx, 'input', pos=0, default=None)
# alpha = get_arg(ctx, 'alpha', pos=1, default=1.0)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
layer = ctx.network.add_activation(input_trt, trt.ActivationType.SELU)
layer.alpha = 1.6732632423543772848170429916717
layer.beta = 1.0507009873554804934193349852946
output._trt = layer.get_output(0)
def convert_roi_pool(ctx):
input = get_arg(ctx, 'input', pos=0, default=None)
boxes = get_arg(ctx, 'boxes', pos=1, default=None)
output_size = get_arg(ctx, 'output_size', pos=2, default=7)
spatial_scale = get_arg(ctx, 'spatial_scale', pos=3, default=1.)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
boxes_trt = trt_(ctx.network, boxes)
plugin = create_roipool_plugin('roi_pool_' + str(id(boxes)),
out_size=output_size,
featmap_strides=[1. / spatial_scale],
roi_scale_factor=-1,
finest_scale=56)
custom_layer = ctx.network.add_plugin_v2(inputs=[boxes_trt, input_trt],
plugin=plugin)
output._trt = custom_layer.get_output(0)
def convert_RoiAlign(ctx):
module = ctx.method_args[0]
input = get_arg(ctx, 'input', pos=1, default=None)
boxes = get_arg(ctx, 'boxes', pos=2, default=None)
output_size = module.output_size
spatial_scale = module.spatial_scale
sampling_ratio = module.sampling_ratio
aligned = module.aligned
old_method_args = ctx.method_args
old_method_kwargs = ctx.method_kwargs
new_method_args = [
input, boxes, output_size, spatial_scale, sampling_ratio, aligned
]
new_method_kwargs = {}
ctx.method_args = new_method_args
ctx.method_kwargs = new_method_kwargs
convert_roi_align(ctx)
ctx.method_args = old_method_args
ctx.method_kwargs = old_method_kwargs
def convert_roi_align(ctx):
input = get_arg(ctx, 'input', pos=0, default=None)
boxes = get_arg(ctx, 'boxes', pos=1, default=None)
output_size = get_arg(ctx, 'output_size', pos=2, default=7)
spatial_scale = get_arg(ctx, 'spatial_scale', pos=3, default=1.)
sampling_ratio = get_arg(ctx, 'sampling_ratio', pos=4, default=-1)
aligned = get_arg(ctx, 'aligned', pos=5, default=False)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
boxes_offset_trt, boxes_trt = trt_(ctx.network, 0.5 / spatial_scale, boxes)
plugin = create_roiextractor_plugin(
'roi_align_' + str(id(boxes)),
out_size=output_size,
sample_num=sampling_ratio,
featmap_strides=[1. / spatial_scale],
roi_scale_factor=1.,
finest_scale=56,
aligned=1 if aligned else 0)
custom_layer = ctx.network.add_plugin_v2(
inputs=[boxes_trt, input_trt], plugin=plugin)
output._trt = custom_layer.get_output(0)
