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 nms_core_symbolic(g, dets, iou_thr, score_thr, max_num):
from torch.onnx.symbolic_opset9 import reshape, unsqueeze, squeeze
from torch.onnx.symbolic_opset10 import _slice
assert 0 <= iou_thr <= 1
multi_bboxes = _slice(g, dets, axes=[1], starts=[0], ends=[4])
# multi_bboxes = unsqueeze(g, multi_bboxes, 0)
multi_bboxes = reshape(g, multi_bboxes, [1, -1, 4])
multi_scores = _slice(g, dets, axes=[1], starts=[4], ends=[5])
multi_scores = reshape(g, multi_scores, [1, 1, -1])
assert max_num > 0
indices = g.op('NonMaxSuppression', multi_bboxes, multi_scores,
g.op('Constant', value_t=torch.LongTensor([max_num])),
g.op('Constant', value_t=torch.FloatTensor([iou_thr])),
g.op('Constant', value_t=torch.FloatTensor([score_thr])))
indices = squeeze(g, _slice(g, indices, axes=[1], starts=[2], ends=[3]), 1)
# Sort indices by score.
scores = reshape(g, multi_scores, [
-1,
])
keeped_scores = g.op('Gather', scores, indices, axis_i=0)
elements_num = sym_help._size_helper(g,
keeped_scores,
dim=g.op('Constant',
value_t=torch.LongTensor(
[0])))
_, order = sym_help._topk_helper(g, keeped_scores, elements_num, dim=0)
indices = g.op('Gather', indices, order, axis_i=0)
return indices
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 squeeze(g, self, dim=None):
if dim is None:
return g.op("Squeeze", self)
dim = sym_help._get_const(dim, 'i', 'dim')
# create 'cond' node (condition is shape[i]==1)
dim_constant = g.op("Constant", value_t=torch.tensor([dim]))
size = sym_help._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()
torch.onnx.utils._add_block(if_node, self, "onnx::Squeeze", axes_i=[dim])
torch.onnx.utils._add_block(if_node, self, "onnx::Identity")
return if_node_outputs
def squeeze(g, self, dim=None):
if dim is None:
return g.op("Squeeze", self)
dim = sym_help._get_const(dim, 'i', 'dim')
input_shape = self.type().sizes()
from torch.onnx.symbolic_helper import _onnx_shape_inference
if input_shape is None or not _onnx_shape_inference:
# 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 = sym_help._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 = torch.onnx.utils._add_block(if_node)
squeeze_ = if_block.op("Squeeze", self, axes_i=[dim])
torch.onnx.utils._add_output_to_block(if_block, squeeze_)
else_block = torch.onnx.utils._add_block(if_node)
identity_ = else_block.op("Identity", self)
torch.onnx.utils._add_output_to_block(else_block, identity_)
return if_node_outputs
# For static input shape
if dim < 0:
dim += self.type().dim()
if input_shape[dim] > 1:
warnings.warn(
"This model contains a squeeze operation on dimension " +
str(dim) + ". The size of " +
"this dimension in the given input is " + str(input_shape[dim]) +
". 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 g.op("Squeeze", self, axes_i=[dim])
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')
loop_condition = g.op("Constant", value_t=torch.tensor(1))
loop_condition = g.op("Cast", loop_condition, to_i=9)
zero = g.op("Constant", value_t=torch.tensor([0]))
indices_len = sym_help._unsqueeze_helper(
g,
sym_help._size_helper(g, indices,
g.op("Constant", value_t=torch.tensor(0))), [0])
if not include_last_offset:
offsets = [offsets, indices_len]
offsets = g.op("Concat", *offsets, axis_i=0)
# Offsets holds the starting index position of each bag. So we create a list of the indices slices (determined by
# offsets) and gather those indices in indices_row. Then we use this subset of indices to gather from embeddings.
# The embeddings output is a loop scan output, so we can avoid creating a sequence and inserting elements in.
offsets_starts = sym_help._slice_helper(g,
offsets,
axes=[0],
starts=[0],
ends=[maxsize],
steps=[1])
offsets_ends = sym_help._slice_helper(g,
offsets,
axes=[0],
starts=[1],
ends=[maxsize],
steps=[1])
loop_len = sym_help._size_helper(g, offsets_ends,
g.op("Constant", value_t=torch.tensor(0)))
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)
indices_start = loop_block.op("Gather",
offsets_starts,
block_input_iter,
axis_i=0)
indices_end = loop_block.op("Gather",
offsets_ends,
block_input_iter,
axis_i=0)
indices_start = sym_help._unsqueeze_helper(loop_block, indices_start, [0])
indices_end = sym_help._unsqueeze_helper(loop_block, indices_end, [0])
indices_row = loop_block.op("Slice", indices, indices_start, indices_end,
zero)
embeddings = loop_block.op("Gather",
embedding_matrix,
indices_row,
axis_i=0)
if not sym_help._is_none(per_sample_weights):
per_sample_weights_row = loop_block.op("Slice", per_sample_weights,
indices_start, indices_end,
zero)
per_sample_weights_row = sym_help._unsqueeze_helper(
loop_block, per_sample_weights_row, [1])
embeddings = loop_block.op("Mul", embeddings, per_sample_weights_row)
if mode == 0:
embeddings = sym_help._reducesum_helper(loop_block,
embeddings,
axes_i=[0],
keepdims_i=0)
elif mode == 1:
embeddings = loop_block.op("ReduceMean",
embeddings,
axes_i=[0],
keepdims_i=0)
else:
embeddings = loop_block.op("ReduceMax",
embeddings,
axes_i=[0],
keepdims_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, embeddings)
# 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 loop.node().output(), None, None, None
def size(g, self, dim=None):
if dim is None:
return g.op("Shape", self)
return sym_help._size_helper(g, self, dim)
def repeat_interleave(g, self, repeats, dim=None):
from torch.onnx.symbolic_opset9 import reshape
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 = reshape(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 = reshape(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 reshape(g, loop_out,
g.op("Constant", value_t=torch.LongTensor(output_sizes)))
def size(g, self, dim):
return sym_help._size_helper(g, self, dim)
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 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 multiclass_nms_core_symbolic(g,
multi_bboxes,
multi_scores,
score_thr,
nms_cfg,
max_num=-1):
from torch.onnx.symbolic_opset9 import reshape, squeeze
from torch.onnx.symbolic_opset10 import _slice
def cast(x, dtype):
return g.op('Cast', x, to_i=sym_help.cast_pytorch_to_onnx[dtype])
def get_size(x, dim):
shape = g.op('Shape', x)
dim = _slice(g, shape, axes=[0], starts=[dim], ends=[dim + 1])
return cast(dim, 'Long')
nms_op_type = nms_cfg.get('type', 'nms')
assert nms_op_type == 'nms'
assert 'iou_thr' in nms_cfg
iou_threshold = nms_cfg['iou_thr']
assert 0 <= iou_threshold <= 1
# Transpose and reshape input tensors to fit ONNX NonMaxSuppression.
multi_bboxes = reshape(g, multi_bboxes, [0, -1, 4])
multi_bboxes = g.op('Transpose', multi_bboxes, perm_i=[1, 0, 2])
batches_num = get_size(multi_bboxes, 0)
spatial_num = get_size(multi_bboxes, 1)
multi_scores = g.op('Transpose', multi_scores, perm_i=[1, 0])
scores_shape = g.op('Concat',
batches_num,
g.op('Constant', value_t=torch.LongTensor([-1])),
spatial_num,
axis_i=0)
multi_scores = reshape(g, multi_scores, scores_shape)
classes_num = get_size(multi_scores, 1)
assert max_num > 0
indices = g.op(
'NonMaxSuppression', multi_bboxes, multi_scores,
g.op('Constant', value_t=torch.LongTensor([max_num])),
g.op('Constant', value_t=torch.FloatTensor([iou_threshold])),
g.op('Constant', value_t=torch.FloatTensor([score_thr])))
# Flatten bboxes and scores.
multi_bboxes_flat = reshape(g, multi_bboxes, [-1, 4])
multi_scores_flat = reshape(g, multi_scores, [
-1,
])
# Flatten indices.
batch_indices = _slice(g, indices, axes=[1], starts=[0], ends=[1])
class_indices = _slice(g, indices, axes=[1], starts=[1], ends=[2])
box_indices = _slice(g, indices, axes=[1], starts=[2], ends=[3])
def add(*args, dtype='Long'):
x = g.op('Add', args[0], args[1])
if dtype is not None:
x = cast(x, dtype)
return x
def mul(*args, dtype='Long'):
x = g.op('Mul', args[0], args[1])
if dtype is not None:
x = cast(x, dtype)
return x
flat_box_indices = add(mul(batch_indices, spatial_num), box_indices)
flat_score_indices = add(
mul(add(mul(batch_indices, classes_num), class_indices), spatial_num),
box_indices)
# Select bboxes.
out_bboxes = reshape(
g, g.op('Gather', multi_bboxes_flat, flat_box_indices, axis_i=0),
[-1, 4])
out_scores = reshape(
g, g.op('Gather', multi_scores_flat, flat_score_indices, axis_i=0),
[-1, 1])
# Having either batch size or number of classes here equal to one is the limitation of implementation.
class_indices = reshape(g, cast(add(class_indices, batch_indices),
'Float'), [-1, 1])
# Combine bboxes, scores and labels into a single tensor.
# This a workaround for a PyTorch bug (feature?),
# limiting ONNX operations to output only single tensor.
out_combined_bboxes = g.op('Concat',
out_bboxes,
out_scores,
class_indices,
axis_i=1)
# Get the top scored bboxes only.
elements_num = sym_help._size_helper(g,
out_scores,
dim=g.op('Constant',
value_t=torch.LongTensor(
[0])))
max_num = g.op('Constant', value_t=torch.LongTensor([max_num]))
if sym_help._export_onnx_opset_version < 12:
kn = g.op('Concat', max_num, elements_num, axis_i=0)
kn = g.op('ReduceMin', kn, keepdims_i=0)
else:
kn = g.op('Min', max_num, elements_num)
_, top_indices = sym_help._topk_helper(g, out_scores, kn, dim=0)
# top_indices = squeeze(g, top_indices, dim=1)
top_indices = reshape(g, top_indices, [
-1,
])
out_combined_bboxes = g.op('Gather',
out_combined_bboxes,
top_indices,
axis_i=0)
return out_combined_bboxes
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")
