wehub-resource-sync
354 lines
11 KiB
Python
354 lines
11 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.
|
|
|
|
|
|
import tvm
|
|
import tvm.script
|
|
import tvm.testing
|
|
from tvm import IRModule, relax
|
|
from tvm.script import relax as R
|
|
|
|
|
|
def _check(
|
|
parsed: relax.Function | IRModule,
|
|
expect: relax.Function | IRModule | None,
|
|
):
|
|
test = parsed.script(show_meta=True)
|
|
roundtrip_mod = tvm.script.from_source(test)
|
|
tvm.ir.assert_structural_equal(parsed, roundtrip_mod)
|
|
if expect:
|
|
tvm.ir.assert_structural_equal(parsed, expect)
|
|
|
|
|
|
def test_all_class_non_max_suppression():
|
|
@R.function
|
|
def foo(
|
|
boxes: R.Tensor((10, 5, 4), "float32"),
|
|
scores: R.Tensor((10, 8, 5), "float32"),
|
|
max_output_boxes_per_class: R.Tensor((), "int64"),
|
|
iou_threshold: R.Tensor((), "float32"),
|
|
score_threshold: R.Tensor((), "float32"),
|
|
) -> R.Tuple(R.Tensor((400, 3), "int64"), R.Tensor((1,), "int64")):
|
|
gv: R.Tuple(R.Tensor((400, 3), "int64"), R.Tensor((1,), "int64")) = (
|
|
R.vision.all_class_non_max_suppression(
|
|
boxes,
|
|
scores,
|
|
max_output_boxes_per_class,
|
|
iou_threshold,
|
|
score_threshold,
|
|
"onnx",
|
|
)
|
|
)
|
|
return gv
|
|
|
|
boxes = relax.Var("boxes", R.Tensor((10, 5, 4), "float32"))
|
|
scores = relax.Var("scores", R.Tensor((10, 8, 5), "float32"))
|
|
max_output_boxes_per_class = relax.Var("max_output_boxes_per_class", R.Tensor((), "int64"))
|
|
iou_threshold = relax.Var("iou_threshold", R.Tensor((), "float32"))
|
|
score_threshold = relax.Var("score_threshold", R.Tensor((), "float32"))
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function(
|
|
"foo", [boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold]
|
|
):
|
|
gv = bb.emit(
|
|
relax.op.vision.all_class_non_max_suppression(
|
|
boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold, "onnx"
|
|
)
|
|
)
|
|
bb.emit_func_output(gv)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_get_valid_counts():
|
|
@R.function
|
|
def foo(
|
|
data: R.Tensor((10, 5, 6), "float32"),
|
|
) -> R.Tuple(
|
|
R.Tensor((10,), "int32"),
|
|
R.Tensor((10, 5, 6), "float32"),
|
|
R.Tensor((10, 5), "int32"),
|
|
):
|
|
gv: R.Tuple(
|
|
R.Tensor((10,), "int32"),
|
|
R.Tensor((10, 5, 6), "float32"),
|
|
R.Tensor((10, 5), "int32"),
|
|
) = R.vision.get_valid_counts(data, score_threshold=0.5, id_index=0, score_index=1)
|
|
return gv
|
|
|
|
data = relax.Var("data", R.Tensor((10, 5, 6), "float32"))
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", [data]):
|
|
gv = bb.emit(
|
|
relax.op.vision.get_valid_counts(data, score_threshold=0.5, id_index=0, score_index=1)
|
|
)
|
|
bb.emit_func_output(gv)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_non_max_suppression_return_indices():
|
|
@R.function
|
|
def foo(
|
|
data: R.Tensor((2, 5, 6), "float32"),
|
|
valid_count: R.Tensor((2,), "int32"),
|
|
indices: R.Tensor((2, 5), "int32"),
|
|
) -> R.Tuple(R.Tensor((2, 5), "int32"), R.Tensor((2, 1), "int32")):
|
|
gv: R.Tuple(R.Tensor((2, 5), "int32"), R.Tensor((2, 1), "int32")) = (
|
|
R.vision.non_max_suppression(
|
|
data,
|
|
valid_count,
|
|
indices,
|
|
max_output_size=-1,
|
|
iou_threshold=0.5,
|
|
force_suppress=False,
|
|
top_k=3,
|
|
coord_start=2,
|
|
score_index=1,
|
|
id_index=0,
|
|
return_indices=True,
|
|
invalid_to_bottom=False,
|
|
soft_nms_sigma=0.0,
|
|
score_threshold=0.0,
|
|
)
|
|
)
|
|
return gv
|
|
|
|
data = relax.Var("data", R.Tensor((2, 5, 6), "float32"))
|
|
valid_count = relax.Var("valid_count", R.Tensor((2,), "int32"))
|
|
indices = relax.Var("indices", R.Tensor((2, 5), "int32"))
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", [data, valid_count, indices]):
|
|
gv = bb.emit(
|
|
relax.op.vision.non_max_suppression(
|
|
data,
|
|
valid_count,
|
|
indices,
|
|
max_output_size=-1,
|
|
iou_threshold=0.5,
|
|
force_suppress=False,
|
|
top_k=3,
|
|
coord_start=2,
|
|
score_index=1,
|
|
id_index=0,
|
|
return_indices=True,
|
|
invalid_to_bottom=False,
|
|
soft_nms_sigma=0.0,
|
|
score_threshold=0.0,
|
|
)
|
|
)
|
|
bb.emit_func_output(gv)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_non_max_suppression_return_indices_soft_nms():
|
|
@R.function
|
|
def foo(
|
|
data: R.Tensor((2, 5, 6), "float32"),
|
|
valid_count: R.Tensor((2,), "int32"),
|
|
indices: R.Tensor((2, 5), "int32"),
|
|
) -> R.Tuple(
|
|
R.Tensor((2, 5, 6), "float32"),
|
|
R.Tensor((2, 5), "int32"),
|
|
R.Tensor((2, 1), "int32"),
|
|
):
|
|
gv: R.Tuple(
|
|
R.Tensor((2, 5, 6), "float32"),
|
|
R.Tensor((2, 5), "int32"),
|
|
R.Tensor((2, 1), "int32"),
|
|
) = R.vision.non_max_suppression(
|
|
data,
|
|
valid_count,
|
|
indices,
|
|
max_output_size=-1,
|
|
iou_threshold=0.5,
|
|
force_suppress=False,
|
|
top_k=3,
|
|
coord_start=2,
|
|
score_index=1,
|
|
id_index=0,
|
|
return_indices=True,
|
|
invalid_to_bottom=False,
|
|
soft_nms_sigma=0.5,
|
|
score_threshold=0.0,
|
|
)
|
|
return gv
|
|
|
|
data = relax.Var("data", R.Tensor((2, 5, 6), "float32"))
|
|
valid_count = relax.Var("valid_count", R.Tensor((2,), "int32"))
|
|
indices = relax.Var("indices", R.Tensor((2, 5), "int32"))
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", [data, valid_count, indices]):
|
|
gv = bb.emit(
|
|
relax.op.vision.non_max_suppression(
|
|
data,
|
|
valid_count,
|
|
indices,
|
|
max_output_size=-1,
|
|
iou_threshold=0.5,
|
|
force_suppress=False,
|
|
top_k=3,
|
|
coord_start=2,
|
|
score_index=1,
|
|
id_index=0,
|
|
return_indices=True,
|
|
invalid_to_bottom=False,
|
|
soft_nms_sigma=0.5,
|
|
score_threshold=0.0,
|
|
)
|
|
)
|
|
bb.emit_func_output(gv)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_non_max_suppression_return_data():
|
|
@R.function
|
|
def foo(
|
|
data: R.Tensor((2, 5, 6), "float32"),
|
|
valid_count: R.Tensor((2,), "int32"),
|
|
indices: R.Tensor((2, 5), "int32"),
|
|
) -> R.Tensor((2, 5, 6), "float32"):
|
|
gv: R.Tensor((2, 5, 6), "float32") = R.vision.non_max_suppression(
|
|
data,
|
|
valid_count,
|
|
indices,
|
|
max_output_size=-1,
|
|
iou_threshold=0.5,
|
|
force_suppress=False,
|
|
top_k=-1,
|
|
coord_start=2,
|
|
score_index=1,
|
|
id_index=0,
|
|
return_indices=False,
|
|
invalid_to_bottom=True,
|
|
soft_nms_sigma=0.0,
|
|
score_threshold=0.0,
|
|
)
|
|
return gv
|
|
|
|
data = relax.Var("data", R.Tensor((2, 5, 6), "float32"))
|
|
valid_count = relax.Var("valid_count", R.Tensor((2,), "int32"))
|
|
indices = relax.Var("indices", R.Tensor((2, 5), "int32"))
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", [data, valid_count, indices]):
|
|
gv = bb.emit(
|
|
relax.op.vision.non_max_suppression(
|
|
data,
|
|
valid_count,
|
|
indices,
|
|
max_output_size=-1,
|
|
iou_threshold=0.5,
|
|
force_suppress=False,
|
|
top_k=-1,
|
|
coord_start=2,
|
|
score_index=1,
|
|
id_index=0,
|
|
return_indices=False,
|
|
invalid_to_bottom=True,
|
|
soft_nms_sigma=0.0,
|
|
score_threshold=0.0,
|
|
)
|
|
)
|
|
bb.emit_func_output(gv)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_multibox_transform_loc():
|
|
@R.function
|
|
def foo(
|
|
cls: R.Tensor((1, 3, 5), "float32"),
|
|
loc: R.Tensor((1, 20), "float32"),
|
|
anc: R.Tensor((1, 5, 4), "float32"),
|
|
) -> R.Tuple(R.Tensor((1, 5, 4), "float32"), R.Tensor((1, 3, 5), "float32")):
|
|
gv: R.Tuple(R.Tensor((1, 5, 4), "float32"), R.Tensor((1, 3, 5), "float32")) = (
|
|
R.vision.multibox_transform_loc(
|
|
cls,
|
|
loc,
|
|
anc,
|
|
clip=False,
|
|
threshold=0.0,
|
|
variances=(1.0, 1.0, 1.0, 1.0),
|
|
keep_background=True,
|
|
)
|
|
)
|
|
return gv
|
|
|
|
cls = relax.Var("cls", R.Tensor((1, 3, 5), "float32"))
|
|
loc = relax.Var("loc", R.Tensor((1, 20), "float32"))
|
|
anc = relax.Var("anc", R.Tensor((1, 5, 4), "float32"))
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", [cls, loc, anc]):
|
|
gv = bb.emit(
|
|
relax.op.vision.multibox_transform_loc(
|
|
cls,
|
|
loc,
|
|
anc,
|
|
clip=False,
|
|
threshold=0.0,
|
|
variances=(1.0, 1.0, 1.0, 1.0),
|
|
keep_background=True,
|
|
)
|
|
)
|
|
bb.emit_func_output(gv)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_roi_align():
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((1, 2, 8, 8), "float32"),
|
|
rois: R.Tensor((2, 5), "float32"),
|
|
) -> R.Tensor((2, 2, 3, 3), "float32"):
|
|
gv: R.Tensor((2, 2, 3, 3), "float32") = R.vision.roi_align(
|
|
x,
|
|
rois,
|
|
pooled_size=(3, 3),
|
|
spatial_scale=1.0,
|
|
sample_ratio=2,
|
|
layout="NCHW",
|
|
mode="avg",
|
|
)
|
|
return gv
|
|
|
|
x = relax.Var("x", R.Tensor((1, 2, 8, 8), "float32"))
|
|
rois = relax.Var("rois", R.Tensor((2, 5), "float32"))
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", [x, rois]):
|
|
gv = bb.emit(
|
|
relax.op.vision.roi_align(
|
|
x, rois, (3, 3), 1.0, sample_ratio=2, layout="NCHW", mode="avg"
|
|
)
|
|
)
|
|
bb.emit_func_output(gv)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|