def symbolic(g, bboxes, scores, iou_threshold, offset):
from ..onnx import is_custom_op_loaded
has_custom_op = is_custom_op_loaded()
if has_custom_op:
return g.op('mmcv::NonMaxSuppression',
bboxes,
scores,
iou_threshold_f=float(iou_threshold),
offset_i=int(offset))
else:
from torch.onnx.symbolic_opset9 import select, squeeze, unsqueeze
boxes = unsqueeze(g, bboxes, 0)
scores = unsqueeze(g, unsqueeze(g, scores, 0), 0)
max_output_per_class = g.op('Constant',
value_t=torch.tensor([sys.maxsize],
dtype=torch.long))
iou_threshold = g.op('Constant',
value_t=torch.tensor([iou_threshold],
dtype=torch.float))
nms_out = g.op('NonMaxSuppression', boxes, scores,
max_output_per_class, iou_threshold)
return squeeze(
g,
select(
g, nms_out, 1,
g.op('Constant',
value_t=torch.tensor([2], dtype=torch.long))), 1)
def symbolic(g, bboxes, scores, iou_threshold, offset):
from ..onnx import is_custom_op_loaded
has_custom_op = is_custom_op_loaded()
# TensorRT nms plugin is aligned with original nms in ONNXRuntime
is_trt_backend = os.environ.get('ONNX_BACKEND') == 'MMCVTensorRT'
if has_custom_op and (not is_trt_backend):
return g.op('mmcv::NonMaxSuppression',
bboxes,
scores,
iou_threshold_f=float(iou_threshold),
offset_i=int(offset))
else:
from torch.onnx.symbolic_opset9 import select, squeeze, unsqueeze
boxes = unsqueeze(g, bboxes, 0)
scores = unsqueeze(g, unsqueeze(g, scores, 0), 0)
max_output_per_class = g.op('Constant',
value_t=torch.tensor([sys.maxsize],
dtype=torch.long))
iou_threshold = g.op('Constant',
value_t=torch.tensor([iou_threshold],
dtype=torch.float))
nms_out = g.op('NonMaxSuppression', boxes, scores,
max_output_per_class, iou_threshold)
return squeeze(
g,
select(
g, nms_out, 1,
g.op('Constant',
value_t=torch.tensor([2], dtype=torch.long))), 1)
def symbolic_multi_label_nms(g, boxes, scores, iou_threshold):
boxes = unsqueeze(g, boxes, 0)
scores = unsqueeze(g, unsqueeze(g, scores, 0), 0)
max_output_per_class = g.op('Constant', value_t=torch.tensor([sys.maxsize], dtype=torch.long))
iou_threshold = g.op('Constant', value_t=torch.tensor([iou_threshold], dtype=torch.float))
nms_out = g.op('NonMaxSuppression', boxes, scores, max_output_per_class, iou_threshold)
return squeeze(g, select(g, nms_out, 1, g.op('Constant', value_t=torch.tensor([2], dtype=torch.long))), 1)
def _interpolate_get_scales_and_mode(g, input, size, scale_factor, mode , align_corners):
from torch.onnx.symbolic_opset9 import unsqueeze
mode = _maybe_get_const(mode, 's')
if 'linear' in mode:
mode = 'linear'
if 'cubic' in mode:
mode = 'cubic'
_interpolate_warning(mode)
align_corners = _maybe_get_const(align_corners, 'b')
if not _is_none(align_corners) and align_corners:
return _unimplemented("interpolate", "align_corners == True")
dim = input.type().dim()
if not _is_none(scale_factor):
scale_factor = _interpolate_get_scales(g, scale_factor, dim)
elif not _is_none(size):
if not _is_packed_list(size):
is_scalar = ((_maybe_get_const(size, 't').dim() == 0))
if is_scalar:
size = unsqueeze(g, size, 0)
size = [size for i in range(dim - 2)]
size = g.op("Concat", *size, axis_i=0)
scale_factor = _interpolate_size_to_scales(g, input, size, dim)
else:
_unimplemented("Both size and scales are None in __interpolate")
return scale_factor, mode
def topk_symbolic(g, self, k, dim, largest, sorted, out=None):
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
if out is not None:
sym_help._unimplemented("TopK",
"Out parameter is not supported for topk")
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))
from torch.onnx.symbolic_opset9 import unsqueeze
k = unsqueeze(g, k, 0)
top_values, top_indices = g.op("TopK", self, k, axis_i=dim, outputs=2)
if not largest:
top_values = reverse(top_values)
top_indices = reverse(top_indices)
return top_values, top_indices
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 symbolic(g, bboxes, scores, iou_threshold, offset):
from torch.onnx.symbolic_opset9 import select, squeeze, unsqueeze
boxes = unsqueeze(g, bboxes, 0)
scores = unsqueeze(g, unsqueeze(g, scores, 0), 0)
max_output_per_class = g.op(
'Constant', value_t=torch.tensor([sys.maxsize], dtype=torch.long))
iou_threshold = g.op(
'Constant',
value_t=torch.tensor([iou_threshold], dtype=torch.float))
nms_out = g.op('NonMaxSuppression', boxes, scores,
max_output_per_class, iou_threshold)
return squeeze(
g,
select(
g, nms_out, 1,
g.op('Constant', value_t=torch.tensor([2], dtype=torch.long))),
1)
def _interpolate_get_scales(g, scale_factor, dim):
from torch.onnx.symbolic_opset9 import unsqueeze
offsets = g.op("Constant", value_t=torch.ones(2, dtype=torch.float32))
if _is_packed_list(scale_factor):
scale_factor = _unpack_list(scale_factor)
scales = []
for dim_scale_factor in scale_factor:
dim_scale_factor = unsqueeze(g, dim_scale_factor, 0)
dim_scale_factor = g.op("Cast", dim_scale_factor, to_i=cast_pytorch_to_onnx["Float"])
scales.append(dim_scale_factor)
else:
scale_factor = unsqueeze(g, scale_factor, 0)
scale_factor = g.op("Cast", scale_factor, to_i=cast_pytorch_to_onnx["Float"])
scales = [scale_factor for i in range(dim - 2)]
scale_factor = g.op("Concat", offsets, *scales, axis_i=0)
return scale_factor
def topk(g, self, k, dim, largest, sorted, out=None):
if out is not None:
_unimplemented("TopK", "Out parameter is not supported for topk")
if not largest:
_unimplemented("TopK", "Ascending TopK is not supported")
k = g.op("Constant", value_t=torch.tensor(k, dtype=torch.int64))
from torch.onnx.symbolic_opset9 import unsqueeze
k = unsqueeze(g, k, 0)
return g.op("TopK", self, k, axis_i=dim, outputs=2)
def _topk_helper(g, input, k, dim, largest=True, sorted=False, out=None):
if out is not None:
_unimplemented("TopK", "Out parameter is not supported")
if _export_onnx_opset_version <= 10:
if not largest:
_unimplemented("TopK", "Ascending is not supported")
return g.op("TopK", input, k_i=k, axis_i=dim, outputs=2)
k = _maybe_get_const(k, 'i')
if not _is_value(k):
k = g.op("Constant", value_t=torch.tensor(k, dtype=torch.int64))
from torch.onnx.symbolic_opset9 import unsqueeze
k = unsqueeze(g, k, 0)
return g.op("TopK", input, k, axis_i=dim, outputs=2)
else:
return g.op("TopK",
input,
k,
axis_i=dim,
largest_i=largest,
sorted_i=sorted,
outputs=2)
def _unsqueeze_helper(g, input, dim):
from torch.onnx.symbolic_opset9 import unsqueeze
return unsqueeze(g, input, dim)
