wehub-resource-sync
195 lines
6.6 KiB
Python
195 lines
6.6 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Default legalization function for vision network related operators."""
|
|
|
|
from tvm import relax, te, tirx, topi
|
|
from tvm.ir import Call
|
|
|
|
from ...block_builder import BlockBuilder
|
|
from ...expr import Expr, TupleGetItem
|
|
from .common import register_legalize
|
|
|
|
|
|
@register_legalize("relax.vision.all_class_non_max_suppression")
|
|
def _all_class_non_max_suppression(block_builder: BlockBuilder, call: Call) -> Expr:
|
|
"""Legalize all_class_non_max_suppression with dynamic output trimming.
|
|
|
|
This implementation uses dynamic_strided_slice to trim the NMS output to only
|
|
contain valid detections, improving memory efficiency and ONNX compatibility.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The legalized NMS result.
|
|
|
|
- For ONNX output format, returns a tuple of
|
|
`(trimmed_indices, num_total_detections)`, where `trimmed_indices`
|
|
contains only valid detection indices.
|
|
- For TensorFlow output format, returns the TOPI result directly to
|
|
preserve the `(selected_indices, selected_scores, num_detections)`
|
|
layout expected by the Relax op.
|
|
"""
|
|
boxes = call.args[0]
|
|
scores = call.args[1]
|
|
max_output_boxes_per_class = call.args[2]
|
|
iou_threshold = call.args[3]
|
|
score_threshold = call.args[4]
|
|
output_format = call.attrs.output_format
|
|
|
|
scores_shape = scores.ty.shape
|
|
if len(scores_shape) == 3:
|
|
_, _, num_boxes = scores_shape
|
|
elif len(scores_shape) == 2:
|
|
_, num_boxes = scores_shape
|
|
else:
|
|
raise ValueError(f"Unexpected scores shape: {scores_shape}")
|
|
|
|
if isinstance(max_output_boxes_per_class, relax.Constant):
|
|
max_boxes_val = int(max_output_boxes_per_class.data.numpy())
|
|
else:
|
|
max_boxes_val = int(num_boxes)
|
|
|
|
# Get NMS result with fixed shape from TOPI
|
|
nms_result = block_builder.call_te(
|
|
topi.vision.all_class_non_max_suppression,
|
|
boxes,
|
|
scores,
|
|
max_boxes_val,
|
|
iou_threshold,
|
|
score_threshold,
|
|
output_format,
|
|
)
|
|
|
|
if output_format == "tensorflow":
|
|
return nms_result
|
|
|
|
selected_indices = block_builder.emit(TupleGetItem(nms_result, 0))
|
|
num_total_detections = block_builder.emit(TupleGetItem(nms_result, 1))
|
|
|
|
# Build slicing parameters using TE to avoid high-level Relax ops during legalization
|
|
def build_begin():
|
|
return te.compute((2,), lambda i: tirx.const(0, "int64"), name="begin")
|
|
|
|
def build_strides():
|
|
return te.compute((2,), lambda i: tirx.const(1, "int64"), name="strides")
|
|
|
|
def build_end(count_tensor):
|
|
# end = [count_tensor[0], 3]
|
|
def compute_end(i):
|
|
return tirx.if_then_else(
|
|
i == 0,
|
|
tirx.Cast("int64", count_tensor[0]),
|
|
tirx.const(3, "int64"),
|
|
)
|
|
|
|
return te.compute((2,), compute_end, name="end")
|
|
|
|
begin = block_builder.call_te(build_begin)
|
|
strides = block_builder.call_te(build_strides)
|
|
end = block_builder.call_te(build_end, num_total_detections)
|
|
|
|
# Apply dynamic strided slice to trim to valid detections only
|
|
trimmed_indices = block_builder.emit(
|
|
relax.op.dynamic_strided_slice(selected_indices, begin, end, strides)
|
|
)
|
|
|
|
# Return trimmed indices along with num_total_detections for compatibility
|
|
return relax.Tuple([trimmed_indices, num_total_detections])
|
|
|
|
|
|
@register_legalize("relax.vision.roi_align")
|
|
def _roi_align(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(
|
|
topi.vision.roi_align,
|
|
call.args[0],
|
|
call.args[1],
|
|
pooled_size=call.attrs.pooled_size,
|
|
spatial_scale=call.attrs.spatial_scale,
|
|
mode=call.attrs.mode,
|
|
sample_ratio=call.attrs.sample_ratio,
|
|
aligned=call.attrs.aligned,
|
|
layout=call.attrs.layout,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.vision.get_valid_counts")
|
|
def _get_valid_counts(block_builder: BlockBuilder, call: Call) -> Expr:
|
|
return block_builder.call_te(
|
|
topi.vision.get_valid_counts,
|
|
call.args[0],
|
|
score_threshold=call.attrs.score_threshold,
|
|
id_index=call.attrs.id_index,
|
|
score_index=call.attrs.score_index,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.vision.non_max_suppression")
|
|
def _non_max_suppression(block_builder: BlockBuilder, call: Call) -> Expr:
|
|
return block_builder.call_te(
|
|
topi.vision.non_max_suppression,
|
|
call.args[0],
|
|
call.args[1],
|
|
call.args[2],
|
|
max_output_size=call.attrs.max_output_size,
|
|
iou_threshold=call.attrs.iou_threshold,
|
|
force_suppress=call.attrs.force_suppress,
|
|
top_k=call.attrs.top_k,
|
|
coord_start=call.attrs.coord_start,
|
|
score_index=call.attrs.score_index,
|
|
id_index=call.attrs.id_index,
|
|
return_indices=call.attrs.return_indices,
|
|
invalid_to_bottom=call.attrs.invalid_to_bottom,
|
|
soft_nms_sigma=call.attrs.soft_nms_sigma,
|
|
score_threshold=call.attrs.score_threshold,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.vision.roi_pool")
|
|
def _roi_pool(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(
|
|
topi.vision.roi_pool,
|
|
call.args[0],
|
|
call.args[1],
|
|
pooled_size=call.attrs.pooled_size,
|
|
spatial_scale=call.attrs.spatial_scale,
|
|
layout=call.attrs.layout,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.vision.multibox_transform_loc")
|
|
def _multibox_transform_loc(bb: BlockBuilder, call: Call) -> Expr:
|
|
variances = tuple(float(x) for x in call.attrs.variances)
|
|
|
|
def _te(cls_pred, loc_pred, anchor):
|
|
return topi.vision.multibox_transform_loc(
|
|
cls_pred,
|
|
loc_pred,
|
|
anchor,
|
|
variances,
|
|
clip=call.attrs.clip,
|
|
threshold=call.attrs.threshold,
|
|
keep_background=call.attrs.keep_background,
|
|
)
|
|
|
|
return bb.call_te(
|
|
_te,
|
|
call.args[0],
|
|
call.args[1],
|
|
call.args[2],
|
|
primfunc_name_hint="multibox_transform_loc",
|
|
)
|