This is an automated email from the ASF dual-hosted git repository.
masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 3ae326cb76 [MetaSchedule] Support Tuning w/ No Cost Model & Fix
Integration Test (#13151)
3ae326cb76 is described below
commit 3ae326cb76ead9de7b5559cb5bcac4be49b0059e
Author: Xiyou Zhou <[email protected]>
AuthorDate: Thu Oct 20 00:13:51 2022 -0700
[MetaSchedule] Support Tuning w/ No Cost Model & Fix Integration Test
(#13151)
Add none for cost model & fix integration test.
---
python/tvm/meta_schedule/cost_model/cost_model.py | 8 +--
tests/python/integration/test_tuning.py | 63 +++++++++++++++--------
2 files changed, 47 insertions(+), 24 deletions(-)
diff --git a/python/tvm/meta_schedule/cost_model/cost_model.py
b/python/tvm/meta_schedule/cost_model/cost_model.py
index 54a4d7a343..f139fcc4e4 100644
--- a/python/tvm/meta_schedule/cost_model/cost_model.py
+++ b/python/tvm/meta_schedule/cost_model/cost_model.py
@@ -106,7 +106,7 @@ class CostModel(Object):
@staticmethod
def create(
- kind: Literal["xgb", "mlp", "random"],
+ kind: Literal["xgb", "mlp", "random", "none"],
*args,
**kwargs,
) -> "CostModel":
@@ -114,8 +114,8 @@ class CostModel(Object):
Parameters
----------
- kind : Literal["xgb", "mlp", "random"]
- The kind of the cost model. Can be "xgb", "mlp", or "random".
+ kind : Literal["xgb", "mlp", "random", "none"]
+ The kind of the cost model. Can be "xgb", "mlp", "random" or
"none".
Returns
-------
@@ -134,6 +134,8 @@ class CostModel(Object):
)
return MLPModel(*args, **kwargs) # type: ignore
+ if kind == "none":
+ return None # no cost model required
raise ValueError(f"Unknown CostModel: {kind}")
diff --git a/tests/python/integration/test_tuning.py
b/tests/python/integration/test_tuning.py
index af51439081..589845b13f 100644
--- a/tests/python/integration/test_tuning.py
+++ b/tests/python/integration/test_tuning.py
@@ -26,6 +26,7 @@ from tvm import meta_schedule as ms
from tvm import relay
from tvm.contrib import graph_executor
from tvm.meta_schedule.testing.relay_workload import get_network
+from tvm.meta_schedule.testing.tune_utils import generate_input_data
from tvm.target.target import Target
logging.basicConfig(
@@ -37,23 +38,21 @@
logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
@pytest.mark.skip("Integration test")
@pytest.mark.parametrize(
- "model_name, input_shape, target, layout",
+ "model_name, input_shape, data_type, target, layout",
[
- ("resnet_18", [1, 3, 224, 224], "llvm --num-cores=16", "NHWC"),
- ("resnet_18", [1, 3, 224, 224], "nvidia/geforce-rtx-3090-ti", "NHWC"),
+ ("resnet_18", [1, 3, 224, 224], "float32", "llvm --num-cores=12",
"NHWC"),
+ ("resnet_18", [1, 3, 224, 224], "float32", "nvidia/geforce-rtx-3070",
"NHWC"),
],
)
def test_meta_schedule_tune_relay(
model_name: str,
input_shape: List[int],
+ data_type: str,
target: str,
layout: Optional[str],
):
dev = tvm.cpu() if str(target).startswith("llvm") else tvm.cuda()
- if model_name.startswith("bert"):
- data = tvm.nd.array(np.random.randint(0, 30521, size=input_shape),
dev) # embedding size
- else:
- data = tvm.nd.array(np.random.randn(*input_shape).astype("float32"),
dev)
+ data = generate_input_data(input_shape, data_type)
mod, params, (input_name, _, _) = get_network(
name=model_name,
@@ -78,22 +77,44 @@ def test_meta_schedule_tune_relay(
params=params,
)
print(profiler.table())
- # Compile without meta-schedule for correctness check
- with tvm.transform.PassContext(opt_level=0):
- rt_mod2 = relay.build(mod, target=target, params=params)
- def get_output(data, lib):
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(input_name, data)
- module.run()
- return module.get_output(0).numpy()
+ def get_output(data, lib, dev):
+ module = graph_executor.GraphModule(lib["default"](dev))
+ module.set_input(input_name, tvm.nd.array(data, device=dev))
+ module.run()
+ return module.get_output(0).numpy()
- # Check correctness
- actual_output = get_output(data, rt_mod1)
- expected_output = get_output(data, rt_mod2)
- assert np.allclose(actual_output, expected_output, rtol=1e-4,
atol=2e-4)
+ # Check correctness
+ actual_output = get_output(data, rt_mod1, dev)
+ print(
+ f"{model_name} finished tuning and running on
{Target(target).kind.name}. "
+ "Running baseline...",
+ flush=True,
+ )
+
+ # Compile without meta-schedule for correctness check
+ baseline_target = "llvm -num-cores=1"
+ with tvm.transform.PassContext(opt_level=0):
+ rt_mod2 = relay.build(mod, target=baseline_target, params=params)
+
+ expected_output = get_output(data, rt_mod2, tvm.cpu())
+ print(
+ f"Basline finished running on {Target(baseline_target).kind.name}. "
+ "Verifying correctness...",
+ flush=True,
+ )
+
+ assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
+ print(
+ f"Correctness verified for {model_name} on
{Target(target).kind.name}.",
+ flush=True,
+ )
if __name__ == """__main__""":
- test_meta_schedule_tune_relay("resnet_18", [1, 3, 224, 224], "llvm
--num-cores=16", "NHWC")
- test_meta_schedule_tune_relay("resnet_18", [1, 3, 224, 224],
"nvidia/geforce-rtx-3090-ti", None)
+ test_meta_schedule_tune_relay(
+ "resnet_18", [1, 3, 224, 224], "float32", "llvm --num-cores=12", "NHWC"
+ )
+ test_meta_schedule_tune_relay(
+ "resnet_18", [1, 3, 224, 224], "float32", "nvidia/geforce-rtx-3070",
None
+ )
