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)