Function test_get_valid_counts_e2e

tests/python/relax/test_op_vision.py:575–620

Run get_valid_counts through legalization and compare with the numpy reference.

()

@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_get_valid_counts_e2e():
    """Run get_valid_counts through legalization and compare with the numpy reference."""

    @tvm.script.ir_module
    class GVCModule:
        @R.function
        def main(
            data: R.Tensor((2, 5, 6), "float32"),
        ) -> R.Tuple(
            R.Tensor((2,), "int32"),
            R.Tensor((2, 5, 6), "float32"),
            R.Tensor((2, 5), "int32"),
        ):
            return R.vision.get_valid_counts(data, score_threshold=0.5, id_index=0, score_index=1)

    data_np = np.array(
        [
            [
                [0.0, 0.95, 0.0, 0.0, 1.0, 1.0],
                [1.0, 0.30, 0.0, 0.0, 1.0, 1.0],
                [-1.0, 0.90, 0.0, 0.0, 1.0, 1.0],
                [2.0, 0.75, 2.0, 2.0, 3.0, 3.0],
                [1.0, 0.10, 4.0, 4.0, 5.0, 5.0],
            ],
            [
                [0.0, 0.55, 0.0, 0.0, 1.0, 1.0],
                [1.0, 0.80, 1.0, 1.0, 2.0, 2.0],
                [2.0, 0.40, 2.0, 2.0, 3.0, 3.0],
                [3.0, 0.60, 3.0, 3.0, 4.0, 4.0],
                [-1.0, 0.95, 5.0, 5.0, 6.0, 6.0],
            ],
        ],
        dtype="float32",
    )
    ref_valid_count, ref_out_data, ref_out_indices = tvm.topi.testing.get_valid_counts_python(
        data_np, score_threshold=0.5, id_index=0, score_index=1
    )

    mod = LegalizeOps()(GVCModule)
    exe = tvm.compile(mod, target="llvm")
    vm = relax.VirtualMachine(exe, tvm.cpu())
    result = vm["main"](tvm.runtime.tensor(data_np, tvm.cpu()))

    tvm.testing.assert_allclose(result[0].numpy(), ref_valid_count)
    tvm.testing.assert_allclose(result[1].numpy(), ref_out_data)
    tvm.testing.assert_allclose(result[2].numpy(), ref_out_indices)


def _prepare_nms_inputs(raw_data: np.ndarray):