def group_norm_symbolic(g, input, num_groups, weight, bias, eps,
cudnn_enabled):
from torch.onnx.symbolic_opset9 import reshape, mul, add, reshape_as
channels_num = input.type().sizes()[1]
if num_groups == channels_num:
output = g.op('InstanceNormalization',
input,
weight,
bias,
epsilon_f=eps)
else:
# Reshape from [n, g * cg, h, w] to [1, n * g, cg * h, w].
x = reshape(g, input, [0, num_groups, -1, 0])
x = reshape(g, x, [1, -1, 0, 0])
# Normalize channel-wise.
x = g.op('MeanVarianceNormalization', x, axes_i=[2, 3])
# Reshape back.
x = reshape_as(g, x, input)
# Apply affine transform.
x = mul(g, x, reshape(g, weight, [1, channels_num, 1, 1]))
output = add(g, x, reshape(g, bias, [1, channels_num, 1, 1]))
return output
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 reverse(x):
from torch.onnx.symbolic_opset9 import reshape, transpose, size
y = transpose(g, x, 0, dim)
shape = g.op("Shape", y)
y = reshape(g, y, [0, 1, -1])
n = size(g, y, g.op("Constant", value_t=torch.LongTensor([0])))
y = g.op("ReverseSequence", y, n, batch_axis_i=1, time_axis_i=0)
y = reshape(g, y, shape)
y = transpose(g, y, 0, dim)
return y
def symbolic(g,
features,
rois,
out_size,
spatial_scale,
sample_num=0,
aligned=True):
batch_indices = reshape(
g,
g.op('Cast',
_slice(g, rois, axes=[1], starts=[0], ends=[1]),
to_i=sym_help.cast_pytorch_to_onnx['Long']), [-1])
bboxes = _slice(g, rois, axes=[1], starts=[1], ends=[5])
if aligned:
scale = sym_help._maybe_get_scalar(spatial_scale)
offset = g.op("Constant",
value_t=torch.tensor(0.5 / scale,
dtype=torch.float32))
bboxes = sub(g, bboxes, offset)
out_h, out_w = _pair(out_size)
return g.op('RoiAlign',
features,
bboxes,
batch_indices,
output_height_i=out_h,
output_width_i=out_w,
sampling_ratio_i=sample_num,
spatial_scale_f=spatial_scale)
def argmin(g, input, dim, keepdim):
if sym_help._is_none(dim):
from torch.onnx.symbolic_opset9 import reshape
flattened = reshape(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 argmin(g, input, dim, keepdim):
if sym_help._is_none(dim):
from torch.onnx.symbolic_opset9 import reshape
flattened = reshape(g, input, (-1,))
return g.op('ArgMin', flattened, axis_i=0, keepdims_i=False, select_last_index_i=True)
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=True)
def reshape(g, input, shape):
if input not in symbolic_helper._quantized_ops:
return opset9.reshape(g, input, shape)
kwargs = {
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
output = g.op("_caffe2::Int8Reshape", input, shape, **kwargs)
symbolic_helper._quantized_ops.add(output)
return output
def reshape(g, input, shape):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import reshape
return reshape(g, input, shape)
kwargs = {
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
output = g.op("_caffe2::Int8Reshape", input, shape, **kwargs)
sym_help._quantized_ops.add(output)
return output
def symbolic(g, features, rois, out_size, spatial_scale, sample_num=0):
batch_indices = reshape(
g,
g.op('Cast',
_slice(g, rois, axes=[1], starts=[0], ends=[1]),
to_i=sym_help.cast_pytorch_to_onnx['Long']), [-1])
bboxes = _slice(g, rois, axes=[1], starts=[1], ends=[5])
out_h, out_w = _pair(out_size)
return g.op('RoiAlign',
features,
bboxes,
batch_indices,
output_height_i=out_h,
output_width_i=out_w,
sampling_ratio_i=sample_num,
spatial_scale_f=spatial_scale)
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 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