def _prepare_onnx_paddings(g, input, pad): if not sym_help._is_packed_list(pad) and sym_help._is_list(pad) and sym_help._is_scalar_list(pad): pad = g.op("ConcatFromSequence", pad, axis_i=0, new_axis_i=1) # The desired order of paddings is # dim_0_begin, dim_1_begin, ... , dim_0_end, ..., dim_n_end. # n is the dimension of input. # Assume zero-dimensions in the beginning, pad the "pad" sequence with zeros in the beginning pad_len = torch.onnx.symbolic_opset9.size(g, pad, g.op("Constant", value_t=torch.tensor([0]))) # Set extension = [0] * (dim * 2 - len(pad)) rank = sym_help._get_tensor_rank(input) if rank is None: rank = g.op("Size", g.op("Shape", input)) else: rank = g.op("Constant", value_t=torch.tensor(rank, dtype=torch.int64)) extension = g.op("Sub", g.op("Mul", rank, g.op("Constant", value_t=torch.tensor(2, dtype=torch.int64))), pad_len) # Concat pad with extension: paddings = [dim_n_begin, dim_n_end, dim_n-1_begin, dim_n-1_end, 0, 0, ... ] # Currently ONNX only supports int64 type for Pad pad = g.op("Cast", pad, to_i=sym_help.cast_pytorch_to_onnx["Long"]) paddings = g.op("Concat", pad, g.op("ConstantOfShape", extension, value_t=torch.tensor([0], dtype=torch.int64)), axis_i=0) # Reshape and reverse order and collate first beginnings and then ends # paddings = [[..., 0, dim_n-1_begin, dim_n_begin], # [..., 0, dim_n-1_end, dim_n_end]] # Reshape back to 1-D paddings = [..., 0, dim_n - 1_begin, dim_n_begin, ..., 0, dim_n - 1_end, dim_n_end] paddings = sym_help._reshape_helper(g, paddings, g.op("Constant", value_t=torch.tensor([-1, 2]))) paddings = g.op("Transpose", torch.onnx.symbolic_opset10.flip(g, paddings, [0]), perm_i=[1, 0]) paddings = sym_help._reshape_helper(g, paddings, g.op("Constant", value_t=torch.tensor([-1]))) padding_c = g.op("Cast", paddings, to_i=sym_help.cast_pytorch_to_onnx["Long"]) return padding_c
def _get_im2col_indices_along_dim(g, input_d, kernel_size_d, dilation_d, padding_d, stride_d): # Input is always 4-D (N, C, H, W) # Calculate indices of sliding blocks along spatial dimension # Slide kernel over input each dim d: # each dimension d ranges from 0 to input[d]+2xpadding[d]-dilation[d]x(kernel_size[d]-1) # with steps = stride blocks_d = g.op("Add", input_d, g.op("Constant", value_t=torch.tensor(padding_d * 2))) blocks_d = g.op( "Sub", blocks_d, g.op("Constant", value_t=torch.tensor(dilation_d * (kernel_size_d - 1)))) # Stride kernel over input and find starting indices along dim d blocks_d_indices = g.op("Range", g.op("Constant", value_t=torch.tensor(0)), blocks_d, g.op("Constant", value_t=torch.tensor(stride_d))) # Apply dilation on kernel and find its indices along dim d kernel_grid = torch.arange(0, kernel_size_d * dilation_d, dilation_d) kernel_grid = g.op("Constant", value_t=kernel_grid.unsqueeze(0)) # Broadcast and add kernel staring positions (indices) with # kernel_grid along dim d, to get block indices along dim d blocks_d_indices = sym_help._unsqueeze_helper(g, blocks_d_indices, [0]) # Reshape to [1, -1] kernel_mask = sym_help._reshape_helper( g, kernel_grid, g.op("Constant", value_t=torch.tensor([-1, 1]))) block_mask = g.op("Add", blocks_d_indices, kernel_mask) return block_mask
def argmin(g, input, dim, keepdim): if sym_help._is_none(dim): flattened = sym_help._reshape_helper(g, input, g.op("Constant", value_t=torch.tensor([-1]))) return g.op("ArgMin", flattened, axis_i=0, keepdims_i=False, select_last_index_i=False) else: dim = _parse_arg(dim, "i") keepdim = _parse_arg(keepdim, "i") return g.op("ArgMin", input, axis_i=dim, keepdims_i=keepdim, select_last_index_i=False)
def linalg_vector_norm(g, self, ord, dim, keepdim, dtype): if ord == 0: if dim is None: self = sym_help._reshape_helper(g, self, g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64))) keepdim = None cond_op = g.op("Not", g.op("Equal", self, g.op("Constant", value_t=torch.LongTensor([0])))) cond_op = g.op("Cast", cond_op, to_i=sym_help.cast_pytorch_to_onnx["Long"]) return sym_help._reducesum_helper(g, cond_op, axes_i=dim, keepdims_i=keepdim) else: return lvn(g, self, ord, dim, keepdim, dtype)
def masked_scatter(g, self, mask, source): from torch.onnx.symbolic_opset9 import nonzero, expand_as, size index = nonzero(g, expand_as(g, mask, self)) # NOTE: source can have more elements than needed. # It could also have arbitrary shape. # This is not supported by ONNX::ScatterND, so we need to flatten and slice source tensor. source = sym_help._reshape_helper(g, source, torch.LongTensor([-1])) source = sym_help._slice_helper(g, source, axes=torch.LongTensor([0]), starts=torch.LongTensor([0]), ends=size(g, index, torch.LongTensor([0])), dynamic_slice=True) return g.op("ScatterND", self, index, source)
def im2col(g, input, kernel_size, dilation, padding, stride): # Input is always 4-D tensor (N, C, H, W) # All other args are int[2] input_h = size(g, input, g.op("Constant", value_t=torch.tensor(2))) input_w = size(g, input, g.op("Constant", value_t=torch.tensor(3))) stride_h, stride_w = stride[0], stride[1] padding_h, padding_w = padding[0], padding[1] dilation_h, dilation_w = dilation[0], dilation[1] kernel_h, kernel_w = kernel_size[0], kernel_size[1] blocks_row_indices = _get_im2col_indices_along_dim(g, input_h, kernel_h, dilation_h, padding_h, stride_h) blocks_col_indices = _get_im2col_indices_along_dim(g, input_w, kernel_w, dilation_w, padding_w, stride_w) output_shape = _get_im2col_output_shape(g, input, kernel_h, kernel_w) padded_input = _get_im2col_padded_input(g, input, padding_h, padding_w) # For a 4D matrix of size (1, 1, 3, 3) as below with kernel_size=2, stride=1, and dilation=1 # [[[[1., 2., 3.,], # [4., 5., 6.,], # [7., 8., 9.,]]]] # First gather indices along rows (dim=2) with blocks_row_indices = [[0,1], [1,2]] to get: # [[[[[1., 2., 3.], # [4., 5., 6.]], # [[4., 5., 6.], # [7., 8., 9.]]]]] # And then gather along cols (dim=4) with blocks_row_indices = [[0,1], [1,2]] to get: # [[[[[[1., 2.], # [4., 5.]], # [[2., 3.], # [5., 6]]], # [[[4., 5.], # [7., 8.]], # [[5., 6.], # [8., 9.]]]]]] # Transpose dims 3 (depth) and 4 (rows), and then reshape to output shape (1, 1, 4, 4) to get: # [[[1., 2., 4., 5.], # [2., 3., 5., 6.], # [4., 5., 7., 8.], # [5., 6., 8., 9.]]] output = g.op("Gather", padded_input, blocks_row_indices, axis_i=2) output = g.op("Gather", output, blocks_col_indices, axis_i=4) output = g.op("Transpose", output, perm_i=[0, 1, 2, 4, 3, 5]) return sym_help._reshape_helper(g, output, output_shape)
def linalg_vector_norm(g, self, ord, dim, keepdim, dtype): if ord == 0: if dim is None: self = symbolic_helper._reshape_helper( g, self, g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64)) ) keepdim = 0 cond_op = g.op( "Not", g.op("Equal", self, g.op("Constant", value_t=torch.LongTensor([0]))) ) cond_op = g.op( "Cast", cond_op, to_i=symbolic_helper.cast_pytorch_to_onnx[self.type().scalarType()], ) return symbolic_helper._reducesum_helper( g, cond_op, axes_i=dim, keepdims_i=keepdim ) else: return opset9.linalg_vector_norm(g, self, ord, dim, keepdim, dtype)
def index_put(g, self, indices_list_value, values, accumulate=False): if sym_help._is_packed_list(indices_list_value): indices_list = sym_help._unpack_list(indices_list_value) else: indices_list = [indices_list_value] if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK: args = [self] + indices_list + [values, accumulate] return g.at("index_put", *args) from torch.onnx.symbolic_opset9 import add, expand accumulate = sym_help._parse_arg(accumulate, "b") if len(indices_list) == 0: return values if len(indices_list) > 1: for idx_ in range(len(indices_list)): if indices_list[idx_].type().scalarType() == "Bool": indices_list[idx_] = g.op("NonZero", indices_list[idx_]) index = indices_list[0] for ind in indices_list[1:]: index = add(g, index, ind) broadcast_index_shape = g.op("Shape", index) indices_list = [ sym_help._unsqueeze_helper( g, expand(g, ind, broadcast_index_shape, None), [-1]) for ind in indices_list ] index = g.op("Concat", *indices_list, axis_i=-1) else: # Replace index_put node with masked_scatter or masked_fill # when inputs to the index_put node contains a single boolean input. # # index_put -> masked_fill # * input index contains single tensor of Bool type (e.g.: %24 <- %23). # * input value contains single element (e.g.: %18). # # Torch IR # %mask : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::clone(%0, %6) # %16 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = # aten::to(%8, %26, %27, %11, %12, %28, %29, %15) # %18 : Float(requires_grad=0, device=cpu) = prim::Constant[value={1}]() # %23 : Bool(8, strides=[1], device=cpu) = aten::view(%16, %22) # %24 : Tensor?[] = prim::ListConstruct(%23) # %25 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = # aten::index_put(%mask, %24, %18, %30) # return (%25) # # # index_put -> masked_scatter # * input index contains single tensor of Bool type (e.g.: %32 <- %31). # * input value contains multiple elements (e.g.: %28). # # Torch IR # %mask : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::clone(%0, %6) # %28 : Float(8, strides=[1], requires_grad=0, device=cpu) # = prim::Constant[value= 1 1 1 1 1 1 1 1 [ CPUFloatType{8} ]]() # %15 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) # = aten::ne(%mask, %some_const) # %23 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) # = aten::to(%15, %34, %35, %18, %19, %36, %37, %22) # %38 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]() # %30 : int[] = prim::Constant[value=[-1]]() # %31 : Bool(8, strides=[1], device=cpu) = aten::view(%23, %30) # %32 : Tensor?[] = prim::ListConstruct(%31) # %33 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) # = aten::index_put(%mask, %32, %28, %38) # return (%33) index = indices_list[0] bool_inp = index if bool_inp.type() is not None and bool_inp.type().scalarType( ) == "Bool": rank = sym_help._get_tensor_rank(values) if rank is not None and rank == 0: from torch.onnx.symbolic_opset9 import masked_fill return masked_fill(g, self, bool_inp, values) return masked_scatter(g, self, bool_inp, values) broadcast_index_shape = g.op("Shape", index) index = sym_help._unsqueeze_helper(g, index, [-1]) sub_data_shape = sym_help._slice_helper(g, g.op("Shape", self), axes=[0], starts=[len(indices_list)], ends=[maxsize]) values_shape = g.op("Concat", broadcast_index_shape, sub_data_shape, axis_i=0) # Check if values is a singular value and expand accordingly rank = sym_help._get_tensor_rank(values) if rank is not None and rank == 0: values = expand(g, values, values_shape, None) values = sym_help._reshape_helper(g, values, values_shape) dtype = self.type().scalarType() if dtype is not None and dtype != values.type().scalarType(): values = g.op("Cast", values, to_i=sym_help.cast_pytorch_to_onnx[dtype]) dtype = sym_help.scalar_type_to_onnx.index( sym_help.cast_pytorch_to_onnx[dtype]) dtype = sym_help.scalar_type_to_pytorch_type[dtype] if accumulate: zeros = g.op("ConstantOfShape", g.op("Shape", self), value_t=torch.tensor([0], dtype=dtype)) result = g.op("ScatterND", zeros, index, values) result = add(g, self, result) else: result = g.op("ScatterND", self, index, values) return result
def repeat_interleave(g, self, repeats, dim=None, output_size=None): input = self final_dim = dim # if dim is None flatten # By default, use the flattened input array, and return a flat output array if sym_help._is_none(dim): input = sym_help._reshape_helper( g, self, g.op("Constant", value_t=torch.tensor([-1]))) dim = 0 else: dim = sym_help._maybe_get_scalar(dim) repeats_dim = sym_help._get_tensor_rank(repeats) repeats_sizes = sym_help._get_tensor_sizes(repeats) input_sizes = sym_help._get_tensor_sizes(input) if repeats_dim is None: raise RuntimeError( "Unsupported: ONNX export of repeat_interleave for unknown " "repeats rank.") if repeats_sizes is None: raise RuntimeError( "Unsupported: ONNX export of repeat_interleave for unknown " "repeats size.") if input_sizes is None: raise RuntimeError( "Unsupported: ONNX export of repeat_interleave for unknown " "input size.") # Handle cases where dim is negative if dim < 0: dim += len(input_sizes) output_sizes = input_sizes.copy() for idx, input_size in enumerate(input_sizes): if input_size is None: output_sizes[idx], input_sizes[idx] = 0, -1 print(output_sizes, input_sizes) cond_dynamic_repeats = (repeats_dim == 1 and repeats_sizes[0] is None) # If input size is dynamic or repeats vector is dynamic if output_sizes[dim] == 0 or cond_dynamic_repeats: reps = sym_help._size_helper(g, input, dim) reps = unsqueeze(g, reps, 0) # Check if repeats vector is a single integer value # or a single dimension tensor with non-dynamic values if repeats_dim == 0 or (repeats_dim == 1 and repeats_sizes[0] == 1): if not sym_help._is_tensor(repeats): repeats = g.op("Constant", value_t=torch.LongTensor(repeats)) repeats = g.op("Expand", repeats, reps) # Check if repeats is dynamic # As repeats is dynamic, we use a where node as a substitute for the if statement # If repests_dim = 1, expand repeats otherwise use original tensor elif cond_dynamic_repeats: repeat_dim = sym_help._size_helper( g, repeats, g.op("Constant", value_t=torch.LongTensor([0]))) repeat_cond = g.op("Equal", repeat_dim, g.op("Constant", value_t=torch.LongTensor([1]))) repeats = where(g, repeat_cond, g.op("Expand", repeats, reps), repeats) # There are cases when the repeats are 1-d tensor with multiple repeats, but dim # provided along one of the dynamic axes provided. A simple example would be # input.shape -> [1, 1, *] where * represents the dynamic axes, and dim = 2 # Now, repeat interleaving can be performed in pytorch when the value of * matches # with the number of elements in repeat, for example if * -> 2, number of repeats # should be 2 as well. else: return torch.onnx.symbolic_opset9.repeat_interleave( g, self, repeats, final_dim) reps_like = g.op("ConstantOfShape", g.op("Shape", repeats), value_t=torch.tensor([1], dtype=torch.long)) r_splits = split(g, repeats, reps_like, 0) i_splits = split(g, input, reps_like, dim) output_sizes[dim], input_sizes[dim] = -1, 1 # Create a loop to iterate over each value along the dimension # and perform individual interleaving using the repeats tensor # Loop is of the following pattern # input (trip_count, cond) # int trip_count = ...; # bool cond = ...; # for (int i=0; i < trip_count && cond; ++i) { # cond = ...; # } # Loop conditions loop_condition = g.op("Constant", value_t=torch.tensor(1)) loop_condition = g.op("Cast", loop_condition, to_i=9) loop_len = reps # Create an empty sequence to store final expansions final_splits = g.op("SequenceEmpty") loop = g.op("Loop", loop_len, loop_condition, final_splits) # Loop inputs loop_block = _add_block(loop.node()) block_input_iter = _add_input_to_block(loop_block) cond = _add_input_to_block(loop_block) final_splits = _add_input_to_block(loop_block) r_split = loop_block.op("SequenceAt", r_splits, block_input_iter) i_split = loop_block.op("SequenceAt", i_splits, block_input_iter) i_split = unsqueeze(loop_block, i_split, dim + 1) r_concat = [ loop_block.op("Constant", value_t=torch.LongTensor(input_sizes[:dim + 1])), r_split, loop_block.op("Constant", value_t=torch.LongTensor(input_sizes[dim + 1:])) ] r_concat = loop_block.op("Concat", *r_concat, axis_i=0) i_split = expand(loop_block, i_split, r_concat, None) i_split = sym_help._reshape_helper( loop_block, i_split, g.op("Constant", value_t=torch.LongTensor(output_sizes))) final_splits = loop_block.op("SequenceInsert", final_splits, i_split) # Loop outputs cond_out = loop_block.op("Cast", loop_condition, to_i=9) _add_output_to_block(loop_block, cond_out) _add_output_to_block(loop_block, final_splits) loop_out = loop.node().output() loop_out = g.op("ConcatFromSequence", loop_out, axis_i=dim) return loop_out
def reshape(g, self, shape): # NOTE: Due to bug in ORT https://github.com/microsoft/onnxruntime/issues/10664 # Reshape export cannot utilize the new allowzero attribute introduced in opset 14. return sym_help._reshape_helper(g, self, shape, allowzero=0)
def repeat_interleave(g, self, repeats, dim=None, output_size=None): input = self final_dim = dim # if dim is None flatten # By default, use the flattened input array, and return a flat output array if sym_help._is_none(dim): input = sym_help._reshape_helper( g, self, g.op("Constant", value_t=torch.tensor([-1]))) dim = 0 else: dim = sym_help._maybe_get_scalar(dim) repeats_dim = sym_help._get_tensor_rank(repeats) repeats_sizes = sym_help._get_tensor_sizes(repeats) input_sizes = sym_help._get_tensor_sizes(input) if repeats_dim is None: raise RuntimeError( "Unsupported: ONNX export of repeat_interleave for unknown " "repeats rank.") if repeats_sizes is None: raise RuntimeError( "Unsupported: ONNX export of repeat_interleave for unknown " "repeats size.") if input_sizes is None: raise RuntimeError( "Unsupported: ONNX export of repeat_interleave for unknown " "input size.") # Handle cases where dim is negative if dim < 0: dim += len(input_sizes) output_sizes = input_sizes.copy() perm_i = [0] for idx, input_size in enumerate(input_sizes): perm_i.append(idx + 1) if input_size is None: output_sizes[idx], input_sizes[idx] = 0, -1 perm_i[0], perm_i[dim] = perm_i[dim], perm_i[0] # Cases when repeats is a single value tensor and dim has unknown input size if (repeats_dim == 0 or (repeats_dim == 1 and repeats_sizes[0] == 1)) and output_sizes[dim] == 0: if not sym_help._is_tensor(repeats): repeats = g.op("Constant", value_t=torch.LongTensor(repeats)) reps = sym_help._size_helper(g, input, dim) reps = unsqueeze(g, reps, 0) repeats = g.op("Expand", repeats, reps) # There are cases when the repeats are 1-d tensor with multiple repeats, but dim # provided along one of the dynamic axes provided. A simple example would be # input.shape -> [1, 1, *] where * represents the dynamic axes, and dim = 2 # Now, repeat interleaving can be performed in pytorch when the value of * matches # with the number of elements in repeat, for example if * -> 2, number of repeats # should be 2 as well. else: return torch.onnx.symbolic_opset9.repeat_interleave( g, self, repeats, final_dim) reps_like = g.op("ConstantOfShape", g.op("Shape", repeats), value_t=torch.tensor([1], dtype=torch.long)) r_splits = split(g, repeats, reps_like, 0) i_splits = split(g, input, reps_like, dim) output_sizes[dim], input_sizes[dim] = -1, 1 # Create a loop to iterate over each value along the dimension # and perform individual interleaving using the repeats tensor # Loop is of the following pattern # input (trip_count, cond) # int trip_count = ...; # bool cond = ...; # for (int i=0; i < trip_count && cond; ++i) { # cond = ...; # } # Loop conditions loop_condition = g.op("Constant", value_t=torch.tensor(1)) loop_condition = g.op("Cast", loop_condition, to_i=9) loop_len = reps loop = g.op("Loop", loop_len, loop_condition) # Loop inputs loop_block = _add_block(loop.node()) block_input_iter = _add_input_to_block(loop_block) cond = _add_input_to_block(loop_block) r_split = loop_block.op("SequenceAt", r_splits, block_input_iter) i_split = loop_block.op("SequenceAt", i_splits, block_input_iter) i_split = unsqueeze(loop_block, i_split, dim + 1) r_concat = [ loop_block.op("Constant", value_t=torch.LongTensor(input_sizes[:dim + 1])), r_split, loop_block.op("Constant", value_t=torch.LongTensor(input_sizes[dim + 1:])) ] r_concat = loop_block.op("Concat", *r_concat, axis_i=0) i_split = expand(loop_block, i_split, r_concat, None) i_split = sym_help._reshape_helper( loop_block, i_split, g.op("Constant", value_t=torch.LongTensor(output_sizes))) # Loop outputs cond_out = loop_block.op("Cast", loop_condition, to_i=9) _add_output_to_block(loop_block, cond_out) _add_output_to_block(loop_block, i_split) loop_out = loop.node().output() # In this loop, the outputs are scan outputs and are concatenated along # the zero'th dimension (by default). In order to avoid this and concatenate # along the dimension provided, some post-processing is required loop_out = g.op("Transpose", loop_out, perm_i=perm_i) return sym_help._reshape_helper( g, loop_out, g.op("Constant", value_t=torch.LongTensor(output_sizes)))
def reshape(g, self, shape): return sym_help._reshape_helper(g, self, shape)