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)