def unsafe_chunk(g, self, chunks, dim, _outputs=None): if _outputs is None: return g.op( "SplitToSequence", self, g.op("Constant", value_t=torch.tensor(1, dtype=torch.long)), axis_i=dim, keepdims_i=0, ) size = symbolic_helper._get_tensor_dim_size(self, dim) if size is None: return symbolic_helper._unimplemented("unsafe_chunk", "unknown dimension size") split_size = (size + chunks - 1) // chunks splits = [split_size] * (size // split_size) leftover = size % split_size if leftover: splits.append(leftover) # TODO: So far we don"t have a module using this method. We"ll keep # this as a constant unless we see a request of dynamics in any # user's modules. splits = g.op("Constant", value_t=torch.tensor(splits, dtype=torch.long)) return g.op("Split", self, splits, axis_i=dim, outputs=_outputs)
def split(g, self, split_size_or_sizes, dim, _outputs=None): if not symbolic_helper._is_split_static(split_size_or_sizes, _outputs): split_out = g.op("SplitToSequence", self, split_size_or_sizes, axis_i=dim) if _outputs is None: return split_out # Convert to multiple slice nodes iff number of splits and number of outputs are statically known. if (symbolic_helper._is_packed_list(split_size_or_sizes) and len(symbolic_helper._unpack_list(split_size_or_sizes)) == _outputs): split_sizes = [ symbolic_helper._unsqueeze_helper(g, v, [0]) for v in symbolic_helper._unpack_list(split_size_or_sizes) ] start = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long)) axis = g.op("Constant", value_t=torch.tensor([dim], dtype=torch.long)) res = [] for i in range(_outputs): end = g.op( "Add", start, split_sizes[i] ) # split_sizes is a list of same length as _outputs res.append(g.op("Slice", self, start, end, axis)) start = end return res return [ g.op( "SequenceAt", split_out, g.op("Constant", value_t=torch.tensor([i], dtype=torch.long)), ) for i in range(_outputs) ] split_val = split_size_or_sizes.node()["value"] if split_val.dim() > 0: return g.op("Split", self, split_size_or_sizes, axis_i=dim, outputs=_outputs) split_size = symbolic_helper._get_const(split_size_or_sizes, "i", "split_size") size = symbolic_helper._get_tensor_dim_size(self, dim) if size is None: if _outputs is not None: size = split_size * _outputs else: raise RuntimeError("Unknown dimension size not supported") splits = [split_size] * (size // split_size) leftover = size % split_size if leftover: splits.append(leftover) splits = g.op("Constant", value_t=torch.tensor(splits)) return g.op("Split", self, splits, axis_i=dim, outputs=_outputs)
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._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK: return g.op("ATen", input, operator_s="unfold", 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 = _add_block(loop.node()) block_input_iter = _add_input_to_block(loop_block) cond = _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) _add_output_to_block(loop_block, cond_out) _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 embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse): output = g.op( "org.pytorch.aten::ATen", weight, indices, padding_idx, scale_grad_by_freq, sparse, operator_s="embedding" ) indices_shape = _get_tensor_sizes(indices) if indices_shape is not None and hasattr(weight.type(), "with_sizes"): output_type = weight.type().with_sizes(indices_shape + [_get_tensor_dim_size(weight, 1)]) output.setType(output_type) return output
def embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse): output = g.op("com.microsoft::ATenOp", weight, indices, padding_idx, scale_grad_by_freq, sparse, name_s='aten::embedding') indices_shape = _get_tensor_sizes(indices) if indices_shape is not None and hasattr(weight.type(), 'with_sizes'): output_type = weight.type().with_sizes( indices_shape + [_get_tensor_dim_size(weight, 1)]) output.setType(output_type) return output
def embedding_bag(g, embedding_matrix, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset, padding_idx): if scale_grad_by_freq and sym_help._training_mode: return sym_help._onnx_unsupported("embedding_bag with scale_grad_by_freq for training mode") if padding_idx is not None and padding_idx >= 0: raise RuntimeError("embedding_bag with padding_idx") from torch.onnx.symbolic_opset9 import select import warnings warnings.warn("Export of embedding_bag with dynamic input/offsets shape is not supported in opset 10. " "Please use opset 11 or higher to export model for dynamic input shape.'") offsets_dim_0 = sym_help._get_tensor_dim_size(offsets, 0) if offsets_dim_0 is not None: if include_last_offset: offset_len = offsets_dim_0 - 1 offsets_extended = offsets else: offset_len = offsets_dim_0 offsets_extended = [offsets, g.op("Constant", value_t=torch.tensor([maxsize]))] offsets_extended = g.op("Concat", *offsets_extended, axis_i=0) list_ = [] for i in range(offset_len): start_ = sym_help._unsqueeze_helper(g, select(g, offsets_extended, torch.tensor(0), torch.tensor(i)), [0]) end_ = sym_help._unsqueeze_helper(g, select(g, offsets_extended, torch.tensor(0), torch.tensor(i + 1)), [0]) axes_ = g.op("Constant", value_t=torch.tensor([0])) indices_row = g.op("Slice", indices, start_, end_, axes_) embeddings = g.op("Gather", embedding_matrix, indices_row) if not sym_help._is_none(per_sample_weights): per_sample_weights_row = g.op("Slice", per_sample_weights, start_, end_, axes_) per_sample_weights_row = sym_help._unsqueeze_helper(g, per_sample_weights_row, [1]) embeddings = g.op("Mul", embeddings, per_sample_weights_row) if mode == 0: embeddings = sym_help._reducesum_helper(g, embeddings, axes_i=[0], keepdims_i=0) elif mode == 1: embeddings = g.op("ReduceMean", embeddings, axes_i=[0], keepdims_i=0) else: embeddings = g.op("ReduceMax", embeddings, axes_i=[0], keepdims_i=0) embeddings = sym_help._unsqueeze_helper(g, embeddings, [0]) list_.append(embeddings) output = g.op("Concat", *list_, axis_i=0) # aten::embedding_bag returns a tuple of 4 elements: output, offset2bag, bag_size, max_indices. # But the last three outputs are not used in torch.nn.EmbeddingBag or torch.nn.functional.embedding_bag. return output, None, None, None else: return sym_help._onnx_unsupported("embedding_bag with unknown shape of offsets for opset 10 is not supported. " "please use opset 11 or higher.")
def squeeze(g, self, dim=None): if dim is None: return g.op("Squeeze", self) # dim as a tensor if not symbolic_helper._is_constant(dim): return symbolic_helper._squeeze_helper(g, self, [dim]) dim = symbolic_helper._get_const(dim, "i", "dim") input_rank = symbolic_helper._get_tensor_rank(self) adjusted_dim = dim if input_rank is not None and dim < 0: adjusted_dim += input_rank dim_size = symbolic_helper._get_tensor_dim_size(self, adjusted_dim) if (dim < 0 and input_rank is None) or dim_size is None: # If onnx shape inference is not on, export always as dynamic. # Because we cannot tell if observed static shape is also static at runtime. # create "cond" node (condition is shape[i]==1) dim_constant = g.op("Constant", value_t=torch.tensor([dim])) size = symbolic_helper._size_helper(g, self, dim_constant) const_one = g.op("Constant", value_t=torch.ones(1, dtype=torch.int64)) cond = g.op("Equal", size, const_one) # create the "If" node and add the "then" and "else" blocks to it. if_node_outputs = g.op("If", cond) if_node = if_node_outputs.node() if_block = utils._add_block(if_node) squeeze_ = symbolic_helper._squeeze_helper(if_block, self, [dim]) utils._add_output_to_block(if_block, squeeze_) else_block = utils._add_block(if_node) identity_ = else_block.op("Identity", self) utils._add_output_to_block(else_block, identity_) return if_node_outputs # For static input shape dim = adjusted_dim if dim_size > 1: warnings.warn( "This model contains a squeeze operation on dimension " + str(dim) + ". The size of " + "this dimension in the given input is " + str(dim_size) + ". The model will " + "be exported without the squeeze node. If the model is intended to be used with dynamic " + "input shapes, please export with dynamic_axes argument." ) return self return symbolic_helper._squeeze_helper(g, self, [dim])
def embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse): custom_attributes_json = ( '{' f'"padding_idx":{str(padding_idx)},' f'"scale_grad_by_freq":{str(scale_grad_by_freq).lower()},' f'"sparse":{str(sparse).lower()}' '}' ) output = g.op("com.microsoft::ATenOp", weight, indices, name_s='aten::embedding', custom_attributes_json_s=custom_attributes_json) # do shape inference and set it via setType indices_shape = _get_tensor_sizes(indices) if indices_shape is not None and hasattr(weight.type(), 'with_sizes'): output_type = weight.type().with_sizes(indices_shape + [_get_tensor_dim_size(weight, 1)]) output.setType(output_type) return output
def embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse): custom_attributes_json = ( "{" f'"padding_idx":{str(padding_idx)},' f'"scale_grad_by_freq":{str(scale_grad_by_freq).lower()},' f'"sparse":{str(sparse).lower()}' "}") output = g.at( "embedding", weight, indices, custom_attributes_json_s=custom_attributes_json, ) # do shape inference and set it via setType indices_shape = _get_tensor_sizes(indices) if indices_shape is not None and hasattr(weight.type(), "with_sizes"): output_type = weight.type().with_sizes( indices_shape + [_get_tensor_dim_size(weight, 1)]) output.setType(output_type) return output
def tensor_split(g, self, indices_or_sections, dim, _outputs=None): axis = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long)) axis = opset11.unsqueeze(g, axis, 0) const_1 = g.op("Constant", value_t=torch.tensor(1, dtype=torch.long)) if symbolic_helper._is_split_static(indices_or_sections, _outputs): split_val = symbolic_helper._node_get(indices_or_sections.node(), "value") if split_val.dim() > 0: start = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long)) res = [] assert _outputs is not None for i in range(_outputs - 1): end = g.op( "Gather", indices_or_sections, g.op("Constant", value_t=torch.tensor([i], dtype=torch.long)), axis_i=0, ) res.append(g.op("Slice", self, start, end, axis)) start = end end = symbolic_helper._size_helper(g, self, axis) res.append(g.op("Slice", self, start, end, axis)) return res split_size = symbolic_helper._get_const( indices_or_sections, "i", "indices_or_sections" ) size = symbolic_helper._get_tensor_dim_size(self, dim) if size is None: if _outputs is not None: size = split_size * _outputs else: raise errors.SymbolicValueError( "Unknown dimension size not supported", self ) min_split_size = size // split_size num_splits_one_extra = size % split_size splits = num_splits_one_extra * [min_split_size + 1] leftover = (split_size - num_splits_one_extra) * [min_split_size] splits = g.op( "Constant", value_t=torch.tensor(splits + leftover, dtype=torch.long) ) return g.op("Split", self, splits, axis_i=dim, outputs=_outputs) if ( symbolic_helper._is_tensor(indices_or_sections) and symbolic_helper._get_tensor_rank(indices_or_sections) == 1 ): loop_len = symbolic_helper._size_helper( g, indices_or_sections, g.op("Constant", value_t=torch.tensor(0)) ) loop_len = opset11.unsqueeze(g, loop_len, 0) loop_condition = g.op("Cast", const_1, to_i=_C_onnx.TensorProtoDataType.BOOL) # To make the first slice in the below loop work, # we pad a zero to the first position so that it will be the initial start of slice. padding_0 = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long)) indices_or_sections = g.op("Concat", padding_0, indices_or_sections, axis_i=0) final_splits = g.op("SequenceEmpty") loop = g.op("Loop", loop_len, loop_condition, final_splits) # Loop inputs loop_block = utils._add_block(loop.node()) block_input_iter = utils._add_input_to_block(loop_block) cond = utils._add_input_to_block(loop_block) final_splits = utils._add_input_to_block(loop_block) start = loop_block.op("Gather", indices_or_sections, block_input_iter, axis_i=0) end = loop_block.op( "Gather", indices_or_sections, loop_block.op("Add", block_input_iter, const_1), axis_i=0, ) slice = loop_block.op("Slice", self, start, end, axis) final_splits = loop_block.op("SequenceInsert", final_splits, slice) # Loop outputs cond_out = loop_block.op("Identity", loop_condition) utils._add_output_to_block(loop_block, cond_out) utils._add_output_to_block(loop_block, final_splits) loop_out = loop.node().output() start = g.op( "Gather", indices_or_sections, g.op("Constant", value_t=torch.tensor(-1, dtype=torch.long)), axis_i=0, ) start = opset11.unsqueeze(g, start, 0) end = symbolic_helper._size_helper(g, self, axis) last_slice = g.op("Slice", self, start, end, axis) return g.op("SequenceInsert", loop_out, last_slice) else: # scalar tensor dim_size = symbolic_helper._size_helper(g, self, axis) min_split_size = g.op("Div", dim_size, indices_or_sections) min_split_size_plus_1 = g.op( "Add", min_split_size, const_1, ) num_splits_one_extra = g.op("Mod", dim_size, indices_or_sections) splits = g.op("Tile", min_split_size_plus_1, num_splits_one_extra) leftover = g.op( "Tile", min_split_size, g.op( "Sub", opset11.unsqueeze(g, indices_or_sections, 0), num_splits_one_extra, ), ) splits = g.op("Concat", splits, leftover, axis_i=0) if _outputs is None: return g.op("SplitToSequence", self, splits, axis_i=dim) return g.op("Split", self, splits, axis_i=dim, outputs=_outputs)
def unfold(g, input, dimension, size, step): const_size = symbolic_helper._maybe_get_const(size, "i") const_step = symbolic_helper._maybe_get_const(step, "i") if not symbolic_helper._is_value( const_size) and not symbolic_helper._is_value(const_step): return opset9.unfold(g, input, dimension, const_size, const_step) if symbolic_helper.is_caffe2_aten_fallback(): return g.at("unfold", input, dimension_i=dimension, size_i=size, step_i=step) sizedim = symbolic_helper._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 = symbolic_helper._size_helper( g, low_indices, g.op("Constant", value_t=torch.tensor(0))) hi_size = symbolic_helper._size_helper( g, hi_indices, g.op("Constant", value_t=torch.tensor(0))) ndim = symbolic_helper._get_tensor_rank(input) assert ndim is not None 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 = utils._add_block(loop.node()) block_input_iter = utils._add_input_to_block(loop_block) cond = 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 = symbolic_helper._unsqueeze_helper(loop_block, starts, [0]) ends = symbolic_helper._unsqueeze_helper(loop_block, ends, [0]) stack = loop_block.op("Slice", input, starts, ends, axes) unsqueeze = symbolic_helper._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) utils._add_output_to_block(loop_block, cond_out) 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 = symbolic_helper._squeeze_helper(g, transpose, [0]) return squeeze else: return symbolic_helper._unimplemented("Unfold", "input size not accessible")