def roll(g, input, shifts, dims): from torch.onnx.symbolic_opset9 import squeeze from packaging import version input_shape = g.op('Shape', input) need_flatten = len(dims) == 0 # If dims is not specified, the tensor will be flattened before # rolling and then restored to the original shape. if need_flatten: resize_shape = input_shape input = g.op('Reshape', input, g.op('Constant', value_t=torch.LongTensor([1, -1]))) input_shape = g.op('Shape', input) dims = [1] for index, dim in enumerate(dims): end_size = sym_help._slice_helper(g, input_shape, axes=[0], ends=[dim + 1], starts=[dim]) shift_size = sym_help._slice_helper(g, shifts, axes=[0], ends=[index + 1], starts=[index]) slice_size = g.op('Sub', end_size, shift_size) # Can not use Mod because tensorrt does not support div_size = g.op('Div', slice_size, end_size) slice_size = g.op('Sub', slice_size, g.op('Mul', end_size, div_size)) if version.parse(torch.__version__) >= version.parse('1.7.0'): # add dim=0 for pytorch 1.9.0 end_size = squeeze(g, end_size, 0) slice_size = squeeze(g, slice_size, 0) else: end_size = g.op('Squeeze', end_size) slice_size = g.op('Squeeze', slice_size) dim = torch.LongTensor([dim]) input_slice0 = sym_help._slice_helper(g, input, axes=dim, starts=torch.LongTensor([0]), ends=slice_size, dynamic_slice=True) input_slice1 = sym_help._slice_helper(g, input, axes=dim, ends=end_size, starts=slice_size, dynamic_slice=True) input = g.op('Concat', input_slice1, input_slice0, axis_i=dim) if need_flatten: input = g.op('Reshape', input, resize_shape) return input
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(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 roi_align(g, input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned): batch_indices = _cast_Long( g, squeeze( g, select( g, rois, 1, g.op('Constant', value_t=torch.tensor([0], dtype=torch.long))), 1), False) rois = select( g, rois, 1, g.op('Constant', value_t=torch.tensor([1, 2, 3, 4], dtype=torch.long))) if aligned: warnings.warn( "ONNX export of ROIAlign with aligned=True does not match PyTorch when using malformed boxes," " ONNX forces ROIs to be 1x1 or larger.") scale = torch.tensor(0.5 / spatial_scale).to(dtype=torch.float) rois = g.op("Sub", rois, scale) return g.op('RoiAlign', input, rois, batch_indices, spatial_scale_f=spatial_scale, output_height_i=pooled_height, output_width_i=pooled_width, sampling_ratio_i=sampling_ratio)
def roi_align(g, input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned): if (aligned): raise RuntimeError( 'Unsupported: ONNX export of roi_align with aligned') batch_indices = _cast_Long( g, squeeze( g, select( g, rois, 1, g.op('Constant', value_t=torch.tensor([0], dtype=torch.long))), 1), False) rois = select( g, rois, 1, g.op('Constant', value_t=torch.tensor([1, 2, 3, 4], dtype=torch.long))) return g.op('RoiAlign', input, rois, batch_indices, spatial_scale_f=spatial_scale, output_height_i=pooled_height, output_width_i=pooled_width, sampling_ratio_i=sampling_ratio)
def symbolic(g, input, rois, output_size, spatial_scale, sampling_ratio, pool_mode, aligned): from torch.onnx.symbolic_opset9 import sub, squeeze from torch.onnx.symbolic_helper import _slice_helper from torch.onnx import TensorProtoDataType # batch_indices = rois[:, 0].long() batch_indices = _slice_helper(g, rois, axes=[1], starts=[0], ends=[1]) batch_indices = squeeze(g, batch_indices, 1) batch_indices = g.op('Cast', batch_indices, to_i=TensorProtoDataType.INT64) # rois = rois[:, 1:] rois = _slice_helper(g, rois, axes=[1], starts=[1], ends=[5]) if aligned: # rois -= 0.5/spatial_scale aligned_offset = g.op('Constant', value_t=torch.tensor([0.5 / spatial_scale], dtype=torch.float32)) rois = sub(g, rois, aligned_offset) # roi align return g.op('RoiAlign', input, rois, batch_indices, output_height_i=output_size[0], output_width_i=output_size[1], spatial_scale_f=spatial_scale, sampling_ratio_i=max(0, sampling_ratio), mode_s=pool_mode)
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 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 symbolic(g, input, rois, output_size, spatial_scale, sampling_ratio, pool_mode, aligned): has_custom_op = False try: import os.path as osp from mmcv.ops import get_onnxruntime_op_path ort_op_path = get_onnxruntime_op_path() has_custom_op = osp.exists(ort_op_path) except ImportError: pass if has_custom_op: return g.op( 'mmcv::MMCVRoiAlign', input, rois, aligned_height_i=output_size[0], aligned_width_i=output_size[1], spatial_scale_f=spatial_scale, sampling_ratio_i=max(0, sampling_ratio), pool_mode_s=pool_mode, aligned_i=aligned) from torch.onnx.symbolic_opset9 import sub, squeeze from torch.onnx.symbolic_helper import _slice_helper from torch.onnx import TensorProtoDataType # batch_indices = rois[:, 0].long() batch_indices = _slice_helper(g, rois, axes=[1], starts=[0], ends=[1]) batch_indices = squeeze(g, batch_indices, 1) batch_indices = g.op( 'Cast', batch_indices, to_i=TensorProtoDataType.INT64) # rois = rois[:, 1:] rois = _slice_helper(g, rois, axes=[1], starts=[1], ends=[5]) if aligned: # rois -= 0.5/spatial_scale aligned_offset = g.op( 'Constant', value_t=torch.tensor([0.5 / spatial_scale], dtype=torch.float32)) rois = sub(g, rois, aligned_offset) # roi align return g.op( 'RoiAlign', input, rois, batch_indices, output_height_i=output_size[0], output_width_i=output_size[1], spatial_scale_f=spatial_scale, sampling_ratio_i=max(0, sampling_ratio), mode_s=pool_mode)
def index(g, self, index): if _is_packed_list(index): indices = sym_help._unpack_list(index) else: indices = [index] if len(indices) == 1: if indices[0].type().scalarType() == "Byte": indices[0] = squeeze(g, nonzero(g, indices[0]), dim=1) return index_select(g, self, 0, indices[0]) else: raise NotImplementedError( "Unsupported aten::index operator with more than 1 indices tensor. " )
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 roi_align(g, input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio): batch_indices = _cast_Long( g, squeeze( g, select( g, rois, 1, g.op('Constant', value_t=torch.tensor([0], dtype=torch.long))), 1), False) rois = select( g, rois, 1, g.op('Constant', value_t=torch.tensor([1, 2, 3, 4], dtype=torch.long))) return g.op('RoiAlign', input, rois, batch_indices, spatial_scale_f=spatial_scale, output_height_i=pooled_height, output_width_i=pooled_width, sampling_ratio_i=sampling_ratio)