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)