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)
