def topk_symbolic(g, self, k, dim, largest, sorted, out=None):
from torch.onnx.symbolic_opset9 import unsqueeze
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
k = sym_help._maybe_get_const(k, 'i')
if not sym_help._is_value(k):
k = g.op("Constant", value_t=torch.tensor(k, dtype=torch.int64))
k = unsqueeze(g, k, 0)
do_reverse = False
if sym_help._export_onnx_opset_version <= 10 and not largest:
do_reverse = True
largest = True
top_values, top_indices = sym_help._topk_helper(g, self, k, dim, largest,
sorted, out)
if sym_help._export_onnx_opset_version <= 10 and do_reverse:
top_values = reverse(top_values)
top_indices = reverse(top_indices)
return top_values, top_indices
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 topk(g, self, k, dim, largest, sorted, out=None):
return sym_help._topk_helper(g,
self,
k,
dim,
largest=largest,
sorted=sorted,
out=out)
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
