def gather(g, self, dim, index, sparse_grad=False): if sym_help._maybe_get_const(sparse_grad, "i"): return _unimplemented("gather", "sparse_grad == True") if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK: return g.op("ATen", self, dim, index, sparse_grad, operator_s="gather") return g.op("GatherElements", self, index, axis_i=dim)
def __interpolate(g, input, size, scale_factor, mode, align_corners, recompute_scale_factor): mode = sym_help._maybe_get_const(mode, 's') if 'linear' in mode: mode = 'linear' if 'cubic' in mode: mode = 'cubic' align_corners = sym_help._maybe_get_const(align_corners, 'b') align_corners = False if not isinstance(align_corners, bool) else align_corners coordinate_transformation_mode = "asymmetric" if mode == "nearest" \ else "align_corners" if align_corners else "pytorch_half_pixel" # roi only takes effect with coordinate_transformation_mode="tf_crop_and_resize" roi = g.op("Constant", value_t=torch.tensor([], dtype=torch.float32)) if not sym_help._is_none(size): input_size = g.op("Shape", input) input_size = sym_help._slice_helper(g, input_size, axes=[0], ends=[2], starts=[0]) # in some cases size is not a packed list but size is a scalar # We need to also verify that (sym_help._maybe_get_const(size, 't').dim() == 0) # but this information is not always available. Try to get the dim, # and if not assume that it is not a scalar. try: is_scalar = not sym_help._is_packed_list(size) and ( (sym_help._maybe_get_const(size, 't').dim() == 0)) except AttributeError: is_scalar = not sym_help._is_packed_list(size) if not is_scalar: warnings.warn( "Cannot verify if the output_size is a scalar " "while exporting interpolate. Assuming that it is not a scalar." ) if is_scalar: if not input.type().dim(): return sym_help._unimplemented( "interpolate (with a scalar output_size)", "missing input shape (try giving an array of output_size values)" ) size = unsqueeze(g, size, 0) size = [size for i in range(input.type().dim() - 2)] size = g.op("Concat", *size, axis_i=0) size = g.op("Cast", size, to_i=sym_help.cast_pytorch_to_onnx['Long']) size = g.op("Concat", input_size, size, axis_i=0) scales = g.op("Constant", value_t=torch.tensor([], dtype=torch.float32)) return g.op( "Resize", input, roi, scales, size, coordinate_transformation_mode_s=coordinate_transformation_mode, cubic_coeff_a_f=-0.75, # only valid when mode="cubic" mode_s=mode, # nearest, linear, or cubic nearest_mode_s="floor") else: # if not sym_help._is_none(scales) if not input.type().dim(): return sym_help._unimplemented("interpolate (with scales)", "missing input shape") scales = sym_help._interpolate_get_scales(g, scale_factor, input.type().dim()) return g.op( "Resize", input, roi, scales, coordinate_transformation_mode_s=coordinate_transformation_mode, cubic_coeff_a_f=-0.75, # only valid when mode="cubic" mode_s=mode, # nearest, linear, or cubic nearest_mode_s="floor") # only valid when mode="nearest"
def pixel_shuffle(g, self, upscale_factor): rank = sym_help._get_tensor_rank(self) if rank is not None and rank != 4: return _unimplemented("pixel_shuffle", "only support 4d input") return g.op("DepthToSpace", self, blocksize_i=upscale_factor, mode_s="CRD")
def gather(g, self, dim, index, sparse_grad=False): if sym_help._maybe_get_const(sparse_grad, "i"): return _unimplemented("gather", "sparse_grad == True") if sym_help.is_caffe2_aten_fallback(): return g.at("gather", self, dim, index, sparse_grad) return g.op("GatherElements", self, index, axis_i=dim)
def pixel_shuffle(g, self, upscale_factor): dims = self.type().sizes() if len(dims) != 4: return _unimplemented("pixel_shuffle", "only support 4d input") return g.op("DepthToSpace", self, blocksize_i=upscale_factor, mode_s="CRD")
def tensordot(g, input_a, input_b, dims_a, dims_b, out=None): if out is not None: _unimplemented("Tensordot", "Out parameter is not supported for tensordot.") dim_count_a = sym_help._get_tensor_rank(input_a) if dim_count_a is None: raise RuntimeError( "Unsupported: ONNX export of tensordot for tensor(input_a) of unknown rank." ) dim_count_b = sym_help._get_tensor_rank(input_b) if dim_count_b is None: raise RuntimeError( "Unsupported: ONNX export of tensordot for tensor(input_b) of unknown rank." ) dims_a = [ (dims_a[i] + dim_count_a) if (dims_a[i] < 0) else dims_a[i] for i in range(len(dims_a)) ] dims_b = [ (dims_b[i] + dim_count_b) if (dims_b[i] < 0) else dims_b[i] for i in range(len(dims_b)) ] left_dims_a = [i for i in range(dim_count_a) if (i not in dims_a)] left_dims_b = [i for i in range(dim_count_b) if (i not in dims_b)] new_input_a = permute(g, input_a, left_dims_a + dims_a) new_input_b = permute(g, input_b, dims_b + left_dims_b) input_shape = g.op("Shape", new_input_a) left_sizes_a = sym_help._slice_helper( g, input_shape, axes=[0], starts=[0], ends=[len(left_dims_a)] ) shape_sizes = [ left_sizes_a, g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)), ] output_a = _reshape_from_tensor(g, new_input_a, shape_sizes) input_shape = g.op("Shape", output_a) slices = sym_help._slice_helper( g, input_shape, axes=[0], starts=[-1], ends=[maxsize] ) shape_sizes = [ g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)), slices, ] output_a = _reshape_from_tensor(g, new_input_a, shape_sizes) input_shape = g.op("Shape", new_input_b) left_sizes_b = sym_help._slice_helper( g, input_shape, axes=[0], starts=[len(dims_b)], ends=[maxsize] ) slices = sym_help._slice_helper( g, input_shape, axes=[0], starts=[0], ends=[len(dims_b)] ) shape_sizes = [ slices, g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)), ] output_b = _reshape_from_tensor(g, new_input_b, shape_sizes) input_shape = g.op("Shape", output_b) slices = sym_help._slice_helper( g, input_shape, axes=[0], starts=[-1], ends=[maxsize] ) shape_sizes = [ g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)), slices, ] output_b = _reshape_from_tensor(g, new_input_b, shape_sizes) output = einsum(g, "ij,jk->ik", g.op("prim::ListConstruct", *[output_a, output_b])) shape_sizes = [left_sizes_a, left_sizes_b] return _reshape_from_tensor(g, output, shape_sizes)
def unfold(g, input, dimension, size, step): const_size = sym_help._maybe_get_const(size, "i") const_step = sym_help._maybe_get_const(step, "i") if not sym_help._is_value(const_size) and not sym_help._is_value(const_step): from torch.onnx.symbolic_opset9 import unfold as _unfold return _unfold(g, input, dimension, const_size, const_step) if sym_help.is_caffe2_aten_fallback(): return g.at("unfold", input, dimension_i=dimension, size_i=size, step_i=step) sizedim = sym_help._get_tensor_dim_size(input, dimension) if sizedim is not None: low_start = g.op("Constant", value_t=torch.tensor(0)) low_end = g.op("Constant", value_t=torch.tensor(sizedim)) hi_end = g.op("Constant", value_t=torch.tensor(sizedim + 1)) low_indices = g.op("Range", low_start, low_end, step) hi_indices = g.op("Range", size, hi_end, step) low_size = sym_help._size_helper( g, low_indices, g.op("Constant", value_t=torch.tensor(0)) ) hi_size = sym_help._size_helper( g, hi_indices, g.op("Constant", value_t=torch.tensor(0)) ) ndim = sym_help._get_tensor_rank(input) perm = list(range(0, ndim)) perm.append(perm.pop(dimension)) unsqueeze_list = [] loop_condition = g.op("Constant", value_t=torch.tensor(1)) loop_condition = g.op("Cast", loop_condition, to_i=9) loop_len = g.op("Min", low_size, hi_size) loop = g.op("Loop", loop_len, loop_condition) loop_block = torch.onnx.utils._add_block(loop.node()) block_input_iter = torch.onnx.utils._add_input_to_block(loop_block) cond = torch.onnx.utils._add_input_to_block(loop_block) starts = loop_block.op("Gather", low_indices, block_input_iter) ends = loop_block.op("Gather", hi_indices, block_input_iter) axes = loop_block.op("Constant", value_t=torch.tensor([2])) starts = sym_help._unsqueeze_helper(loop_block, starts, [0]) ends = sym_help._unsqueeze_helper(loop_block, ends, [0]) stack = loop_block.op("Slice", input, starts, ends, axes) unsqueeze = sym_help._unsqueeze_helper( loop_block, loop_block.op("Transpose", stack, perm_i=perm), [dimension] ) unsqueeze_list.append(unsqueeze) concat = loop_block.op("Concat", *unsqueeze_list, axis_i=0) cond_out = loop_block.op("Cast", loop_condition, to_i=9) torch.onnx.utils._add_output_to_block(loop_block, cond_out) torch.onnx.utils._add_output_to_block(loop_block, concat) loop_output = loop.node().output() perm = [0, 1, 2, 3, 4] perm[0], perm[dimension + 1] = perm[dimension + 1], perm[0] transpose = g.op("Transpose", loop_output, perm_i=perm) squeeze = sym_help._squeeze_helper(g, transpose, [0]) return squeeze else: return _unimplemented("Unfold", "input size not accessible")
def reduce_dim(g, self, dim, keepdim, dtype): if dtype.node().kind() != 'prim::Constant': return _unimplemented(name, "dtype") return symbolic(g, self, dim, keepdim)
def pixel_unshuffle(g, self, downscale_factor): rank = sym_help._get_tensor_rank(self) if rank is not None and rank != 4: return _unimplemented("pixel_unshuffle", "only support 4d input") return g.op("SpaceToDepth", self, blocksize_i=downscale_factor)
def diagonal(g, self, offset, dim1, dim2): dim1_size = size(g, self, dim=g.op("Constant", value_t=torch.LongTensor([dim1]))) dim2_size = size(g, self, dim=g.op("Constant", value_t=torch.LongTensor([dim2]))) # Create appropriate mask mask_shape = g.op("Concat", dim1_size, dim2_size, axis_i=0) mask = zeros(g, mask_shape, None, None, None) mask = g.op("EyeLike", mask, k_i=offset) # dim1 and dim2 appended as a dimension at the end of the shape rank = sym_help._get_tensor_rank(self) if rank is not None: axes = list(range(rank)) axes.remove(dim1) axes.remove(dim2) self = g.op("Transpose", self, perm_i=axes + [dim1, dim2]) else: return _unimplemented("diagonal", "unknown input rank") # Multiply input and mask to calculate values along diagonal # The mask consists of one values where diagonal values are to be calculated # For example: # [[1.1, 1.2, 1.3], * [[1, 0, 0] = [[1.1, 0, 0], # [2.1, 2.2, 2.3], [0, 1, 0] [0, 2.2, 0], # [3.1, 3.2, 3.3]] [0, 0, 1]] [0, 0, 3.3]] result = g.op("Mul", self, mask) result = sym_help._reducesum_helper(g, result, axes_i=[-1], keepdims_i=0) # Calculate gather indices based on offset and dims # If offset is greater than zero, set offset to zero as this aids in # calculation of selection window offset_op = g.op("Constant", value_t=torch.LongTensor([offset])) if offset >= 0: diag_size = g.op( "Max", g.op("Min", dim1_size, g.op("Sub", dim2_size, offset_op)), g.op("Constant", value_t=torch.LongTensor([0]))) offset = 0 else: diag_size = g.op( "Max", g.op("Min", g.op("Add", dim1_size, offset_op), dim2_size), g.op("Constant", value_t=torch.LongTensor([0]))) diag_size = g.op("Concat", diag_size, axis_i=0) # Calculate which diagonal values to select # For example, in cases with offsets: # [[0, 1.1, 0] # [0, 0, 2.2]] # we need to select the last two columns, so we create a tensor # with all columns that are to be selected # So in this example, it is [1, 2] select_window_ones_fill = ones(g, diag_size, 4, None, None) select_window = g.op("CumSum", select_window_ones_fill, g.op("Constant", value_t=torch.LongTensor([0]))) select_window = g.op( "Add", select_window, g.op("Constant", value_t=torch.LongTensor([abs(offset) - 1]))) gather_shape = [ size(g, result, dim=g.op("Constant", value_t=torch.LongTensor([axis]))) for axis in list(range(rank))[:-2] ] gather_shape.append(diag_size) gather_shape = g.op("Concat", *gather_shape, axis_i=0) gather_indices = zeros(g, gather_shape, 4, None, None) # There might be cases where offset value is greater than number of rows/columns # and might cause the diagonal to overrun and as a result of this, diag_size would be zero. # For example, if # offset = 9, dim1_size = 2 (columns), dim2_size = 4 (rows) # diag_size = max(min(2, (4-9)), 0) = 0, based on calculation above # Cases with diagonal overrun always result in diag_size = max(0, -ve value) = 0 # In cases without diagonal overrun, we select the appropriate rows/columns along which we # are calculating diagonal values. In cases with diagonal overrun, we return a tensor which has # the dimension of the row/column where overrun occurred as 0-dim, as we are essentially # returning an empty tensor overrun_cond = g.op( "Not", g.op("Equal", diag_size, g.op("Constant", value_t=torch.tensor(0, dtype=torch.int64)))) if_op = g.op("If", overrun_cond) if_node = if_op.node() if_block = _add_block(if_node) gather_indices_if_block = if_block.op("Add", gather_indices, select_window) gather_indices_if_block = sym_help._unsqueeze_helper( if_block, gather_indices_if_block, [rank - 1]) final_non_overrun_ = if_block.op("GatherND", result, gather_indices_if_block, batch_dims_i=rank - 2) _add_output_to_block(if_block, final_non_overrun_) else_block = _add_block(if_node) final_overrun_ = zeros(else_block, gather_shape, 6, None, None) _add_output_to_block(else_block, final_overrun_) return if_op