def convert_expand(ctx):
input = ctx.method_args[0]
if isinstance(ctx.method_args[1], int):
sizes = ctx.method_args[1:]
else:
sizes = ctx.method_args[1]
output = ctx.method_return
input_trt = trt_(ctx.network, input)
dim = len(sizes)
# unsqueeze if necessary
if len(input_trt.shape) < dim:
input_trt = _unsqueeze_input(ctx, input_trt, dim)
input_shape_trt = tensor_trt_get_shape_trt(ctx.network, input_trt)
# compute output shape
output_shape_trt = []
for i, s in enumerate(sizes):
if s > 0:
output_shape_trt.append(trt_(ctx.network, s))
else:
output_shape_trt.append(
slice_shape_trt(ctx.network, input_shape_trt, i, 1))
output_shape_trt = ctx.network.add_concatenation(
output_shape_trt).get_output(0)
# convert repeat
output_trt = _convert_repeat_impl(ctx, input_trt, output_shape_trt)
output._trt = output_trt
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_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)
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_meshgrid(ctx):
input_list = ctx.method_args
output = ctx.method_return
num_inputs = len(input_list)
input_trt_list = [
trt_(ctx.network, input_tensor) for input_tensor in input_list
]
input_shape_trt_list = [
tensor_trt_get_shape_trt(ctx.network, input_trt)
for input_trt in input_trt_list
]
output_shape_trt = ctx.network.add_concatenation(
input_shape_trt_list).get_output(0)
one_trt = trt_(ctx.network, torch.ones(1, dtype=torch.int32))
for index, input_trt in enumerate(input_trt_list):
shuffle_shape_trt = [one_trt] * index
shuffle_shape_trt += [input_shape_trt_list[index]]
shuffle_shape_trt += [one_trt] * (num_inputs - 1 - index)
shuffle_shape_trt = \
ctx.network.add_concatenation(shuffle_shape_trt).get_output(0)
layer = ctx.network.add_shuffle(input_trt)
layer.set_input(1, shuffle_shape_trt)
input_trt_list[index] = layer.get_output(0)
for input_trt, out in zip(input_trt_list, output):
out._trt = _convert_repeat_impl(ctx, input_trt, output_shape_trt)
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_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_BatchNorm1d(ctx):
module = ctx.method_args[0]
input = ctx.method_args[1]
input_trt = trt_(ctx.network, input)
output = ctx.method_return
scale = module.weight.detach().cpu().numpy() / np.sqrt(
module.running_var.detach().cpu().numpy() + module.eps)
bias = module.bias.detach().cpu().numpy(
) - module.running_mean.detach().cpu().numpy() * scale
power = np.ones_like(scale)
# reshape to 2D
input_shape_trt = ctx.network.add_shape(input_trt).get_output(0)
one_trt = trt_(ctx.network,
torch.tensor([1], dtype=torch.int32).to(input.device))
if len(input.shape) == 2:
new_input_shape_trt = ctx.network.add_concatenation(
[input_shape_trt, one_trt, one_trt]).get_output(0)
else:
new_input_shape_trt = ctx.network.add_concatenation(
[input_shape_trt, one_trt]).get_output(0)
layer = ctx.network.add_shuffle(input_trt)
layer.set_input(1, new_input_shape_trt)
layer = ctx.network.add_scale(
layer.get_output(0), trt.ScaleMode.CHANNEL, bias, scale, power)
# reshape back to 1D
conv_out_trt = layer.get_output(0)
layer = ctx.network.add_shuffle(conv_out_trt)
layer.set_input(1, input_shape_trt)
output._trt = layer.get_output(0)
def convert_new_ones(ctx):
input = ctx.method_args[0]
size = ctx.method_args[1]
if isinstance(size, int):
size = ctx.method_args[1:]
dtype = input.dtype
if 'dtype' in ctx.method_kwargs:
dtype = ctx.method_kwargs['dtype']
output = ctx.method_return
if isinstance(size, int):
size = (size, )
# check const
is_const = True
for s in size:
if hasattr(s, '_trt'):
is_const = False
break
if is_const:
# create const value
output_trt = trt_(ctx.network, output)
else:
# create fill
trt_size = []
for s in size:
if hasattr(s, '_trt'):
trt_size.append(s._trt)
else:
trt_size.append(trt_(ctx.network, s))
trt_size = ctx.network.add_concatenation(trt_size).get_output(0)
layer = ctx.network.add_fill(size, trt.FillOperation.RANDOM_UNIFORM)
layer.set_input(0, trt_size)
layer.set_input(1, trt_(ctx.network, input.new_tensor(1)))
layer.set_input(2, trt_(ctx.network, input.new_tensor(1)))
output_trt = layer.get_output(0)
data_type = None
if dtype == torch.float32:
data_type = trt.DataType.FLOAT
elif dtype == torch.int32 or dtype == torch.long:
data_type = trt.DataType.INT32
elif dtype == torch.bool:
data_type = trt.DataType.BOOL
else:
print('unsupported convert type:{}'.format(dtype))
if data_type is not None:
layer = ctx.network.add_identity(output_trt)
layer.set_output_type(0, data_type)
output_trt = layer.get_output(0)
output._trt = output_trt
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_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_type_as(ctx):
input = ctx.method_args[0]
other = ctx.method_args[1]
output = ctx.method_return
input_trt = trt_(ctx.network, input)
other_trt = trt_(ctx.network, other)
layer = ctx.network.add_identity(input_trt)
layer.set_output_type(0, other_trt.dtype)
output._trt = layer.get_output(0)
output._trt.shape # trick to enable type cast
def convert_view_as(ctx):
input = ctx.method_args[0]
other = get_arg(ctx, 'other', pos=1, default=None)
input_trt = trt_(ctx.network, input)
other_trt = trt_(ctx.network, other)
output = ctx.method_return
shape_trt = ctx.network.add_shape(other_trt).get_output(0)
layer = ctx.network.add_shuffle(input_trt)
layer.set_input(1, shape_trt)
output._trt = layer.get_output(0)
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_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_view(ctx):
input = ctx.method_args[0]
size = get_arg(ctx, 'shape', pos=1, default=[])
if isinstance(size, int):
size = tuple(ctx.method_args[1:])
input_trt = trt_(ctx.network, input)
output = ctx.method_return
if input.dtype == torch.bool:
input_trt = trt_cast(ctx.network, input_trt, torch.int32)
# check if there are shape tensor
is_shape_tensor = False
for s in size:
if isinstance(s, IntWarper):
is_shape_tensor = True
break
# negative shape might cause overflow, forbid for now
for s in size:
if s < 0:
is_shape_tensor = True
break
# compute shape tensor
if is_shape_tensor:
shape_trt = []
for idx, s in enumerate(size):
if isinstance(s, IntWarper):
shape_trt.append(s._trt)
else:
const_shape_trt = trt_(
ctx.network, input.new_tensor([s], dtype=torch.int32))
shape_trt.append(const_shape_trt)
shape_trt = ctx.network.add_concatenation(shape_trt).get_output(0)
layer = ctx.network.add_shuffle(input_trt)
if is_shape_tensor:
layer.set_input(1, shape_trt)
else:
layer.reshape_dims = output.shape
output_trt = layer.get_output(0)
if input.dtype == torch.bool:
output_trt = trt_cast(ctx.network, output_trt, torch.bool)
output._trt = output_trt
def __add_clamp(network, trt_input, val, op):
# create TensorRT constant for minimum value
val_shape = (1, ) * len(trt_input.shape) # broadcast all dimensions
if isinstance(val, IntWarper):
val_trt = val._trt
if version.parse(trt.__version__) < version.parse('8'):
# convert type
layer = network.add_identity(val_trt)
layer.set_output_type(0, trt_input.dtype)
val_trt = layer.get_output(0)
# convert 2 / to prevent warning, might remove in future version
layer = network.add_elementwise(
val_trt, trt_(network, torch.zeros(
(1, ), dtype=torch.float32)), trt.ElementWiseOperation.SUM)
layer.set_output_type(0, trt_input.dtype)
val_trt = layer.get_output(0)
else:
torch_type = torch_dtype_from_trt(val_trt.dtype)
# convert 2 / to prevent warning, might remove in future version
layer = network.add_elementwise(
val_trt, trt_(network, torch.zeros((1, ), dtype=torch_type)),
trt.ElementWiseOperation.SUM)
layer.set_output_type(0, trt_input.dtype)
val_trt = layer.get_output(0)
# convert type
layer = network.add_identity(val_trt)
layer.set_output_type(0, trt_input.dtype)
val_trt = layer.get_output(0)
# reshape
layer = network.add_shuffle(val_trt)
layer.reshape_dims = val_shape
val_trt = layer.get_output(0)
else:
# val_shape = (1, ) * len(trt_input.shape) # broadcast all dimensions
val_tensor = val * torch.ones(
val_shape, dtype=torch_dtype_from_trt(
trt_input.dtype)).cpu().numpy()
layer = network.add_constant(val_shape, val_tensor)
val_trt = layer.get_output(0)
layer = network.add_elementwise(trt_input, val_trt, op)
return layer
def convert_expand_as(ctx):
input = ctx.method_args[0]
other = get_arg(ctx, 'other', pos=1, default=None)
input_trt = trt_(ctx.network, input)
other_trt = trt_(ctx.network, other)
output = ctx.method_return
# compute output shape
output_shape_trt = tensor_trt_get_shape_trt(ctx.network, other_trt)
# convert repeat
output_trt = _convert_repeat_impl(ctx, input_trt, output_shape_trt)
output._trt = output_trt
def convert_Conv1d(ctx):
module = ctx.method_args[0]
input = ctx.method_args[1]
input_trt = trt_(ctx.network, input)
output = ctx.method_return
kernel_size = (module.kernel_size[0], 1)
stride = (module.stride[0], 1)
padding = (module.padding[0], 0)
dilation = (module.dilation[0], 1)
kernel = module.weight.detach().cpu().numpy()[..., None]
bias = trt.Weights(torch_dtype_to_trt(module.weight.dtype))
if module.bias is not None:
bias = module.bias.detach().cpu().numpy()
# reshape to 2D
input_shape_trt = ctx.network.add_shape(input_trt).get_output(0)
one_trt = trt_(ctx.network,
torch.tensor([1], dtype=torch.int32).to(input.device))
new_input_shape_trt = ctx.network.add_concatenation(
[input_shape_trt, one_trt]).get_output(0)
layer = ctx.network.add_shuffle(input_trt)
layer.set_input(1, new_input_shape_trt)
layer = ctx.network.add_convolution(input=layer.get_output(0),
num_output_maps=module.out_channels,
kernel_shape=kernel_size,
kernel=kernel,
bias=bias)
layer.stride = stride
layer.padding = padding
layer.dilation = dilation
if module.groups is not None:
layer.num_groups = module.groups
# reshape back to 1D
conv_out_trt = layer.get_output(0)
out_shape_trt = ctx.network.add_shape(conv_out_trt).get_output(0)
new_out_shape_trt = ctx.network.add_slice(out_shape_trt, [0], [3],
[1]).get_output(0)
layer = ctx.network.add_shuffle(conv_out_trt)
layer.set_input(1, new_out_shape_trt)
output._trt = layer.get_output(0)
def convert_sigmoid(ctx):
input = ctx.method_args[0]
input_trt = trt_(ctx.network, input)
output = ctx.method_return
layer = ctx.network.add_activation(input_trt, trt.ActivationType.SIGMOID)
output._trt = layer.get_output(0)
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_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_ReLU(ctx):
input = ctx.method_args[1]
input_trt = trt_(ctx.network, input)
output = ctx.method_return
layer = ctx.network.add_activation(input=input_trt,
type=trt.ActivationType.RELU)
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_add(ctx):
input_a = ctx.method_args[0]
input_b = ctx.method_args[1]
input_a_trt, input_b_trt = trt_(ctx.network, input_a, input_b)
output = ctx.method_return
layer = ctx.network.add_elementwise(input_a_trt, input_b_trt,
trt.ElementWiseOperation.SUM)
output._trt = layer.get_output(0)
def convert_addcmul(ctx):
tensor0 = ctx.method_args[0]
value = 1
next_tensor_offset = 0
if len(ctx.method_args) == 4:
value = ctx.method_args[1]
next_tensor_offset = 1
if 'value' in ctx.method_kwargs:
value = ctx.method_kwargs['value']
tensor1 = ctx.method_args[1 + next_tensor_offset]
tensor2 = ctx.method_args[2 + next_tensor_offset]
input0_trt, input1_trt, input2_trt = trt_(ctx.network, tensor0, tensor1,
tensor2)
output = ctx.method_return
output_mul_trt = ctx.network.add_elementwise(
input1_trt, input2_trt, trt.ElementWiseOperation.PROD).get_output(0)
if value != 1 or value != 1.:
shift = np.zeros([1], np.float32)
scale = np.array([value], np.float32)
if len(tensor0.shape) < 4:
input_shape_trt = tensor_trt_get_shape_trt(ctx.network, input0_trt)
add_dim = 4 - len(tensor0.shape)
add_trt = trt_(ctx.network, torch.ones([add_dim],
dtype=torch.int32))
new_input_shape_trt = ctx.network.add_concatenation(
[add_trt, input_shape_trt]).get_output(0)
layer = ctx.network.add_shuffle(output_mul_trt)
layer.set_input(1, new_input_shape_trt)
output_mul_trt = layer.get_output(0)
output_mul_trt = ctx.network.add_scale(output_mul_trt,
trt.ScaleMode.UNIFORM, shift,
scale).get_output(0)
if len(tensor0.shape) < 4:
layer = ctx.network.add_shuffle(output_mul_trt)
layer.set_input(1, input_shape_trt)
output_mul_trt = layer.get_output(0)
output_trt = ctx.network.add_elementwise(
input0_trt, output_mul_trt, trt.ElementWiseOperation.SUM).get_output(0)
output._trt = output_trt
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_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)
