Function test_roi_align

tests/python/relax/test_tvmscript_parser_op_vision.py:320–349

()

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__":