Add simple testing model.

python/tvm/driver/tvmc/autotuner.py

@@ -374,8 +374,7 @@ def tune_model(
         The path to the produced tuning log file.
     """
     target, extra_targets = common.target_from_cli(target)
-    if target_host is not None:
-        target, target_host = Target.check_and_update_host_consist(target, target_host)
+    target, target_host = Target.check_and_update_host_consist(target, target_host)
     # TODO(jwfromm) Remove this deepcopy once AlterOpLayout bug that mutates source
     # model is fixed. For now, creating a clone avoids the issue.
     mod = deepcopy(tvmc_model.mod)

python/tvm/driver/tvmc/compiler.py

@@ -200,8 +200,7 @@ def compile_model(
         mod = common.convert_graph_layout(mod, desired_layout)
 
     tvm_target, extra_targets = common.target_from_cli(target)
-    if target_host is not None:
-        tvm_target, target_host = Target.check_and_update_host_consist(tvm_target, target_host)
+    tvm_target, target_host = Target.check_and_update_host_consist(tvm_target, target_host)
 
     for codegen_from_cli in extra_targets:
         codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"])

tests/python/driver/tvmc/conftest.py

@@ -102,6 +102,29 @@ def keras_resnet50(tmpdir_factory):
     return model_file_name
 
 
+@pytest.fixture(scope="session")
+def keras_simple(tmpdir_factory):
+    try:
+        from tensorflow import keras
+    except ImportError:
+        # not all environments provide TensorFlow, so skip this fixture
+        # if that is that case.
+        return ""
+
+    model_file_name = "{}/{}".format(tmpdir_factory.mktemp("data"), "simple_conv.h5")
+    model = keras.Sequential(
+        [
+            keras.layers.InputLayer(input_shape=[32, 32, 3], batch_size=1),
+            keras.layers.Conv2D(8, kernel_size=(3, 3)),
+            keras.layers.Flatten(),
+            keras.layers.Dense(64),
+        ]
+    )
+    model.save(model_file_name)
+
+    return model_file_name
+
+
 @pytest.fixture(scope="session")
 def pytorch_resnet18(tmpdir_factory):
     try:

tests/python/driver/tvmc/test_autoscheduler.py

@@ -44,7 +44,7 @@ def _autoscheduler_test_helper(model, tmpdir_name, early_stopping=1, prior_recor
         prior_records=prior_records,
         early_stopping=early_stopping,
         enable_autoscheduler=True,
-        trials=5,
+        trials=2,
         hardware_params=hardware_params,
     )
 
@@ -59,36 +59,36 @@ def _autoscheduler_test_helper(model, tmpdir_name, early_stopping=1, prior_recor
     return log_file
 
 
-def test_get_tuning_tasks(onnx_mnist):
-    pytest.importorskip("onnx")
+def test_get_tuning_tasks(keras_simple):
+    pytest.importorskip("tensorflow")
 
-    tasks, weights = _get_tasks(onnx_mnist)
+    tasks, weights = _get_tasks(keras_simple)
     expected_task_type = auto_scheduler.SearchTask
 
     assert type(tasks) is list
     assert len(tasks) > 0
     assert all([type(x) is expected_task_type for x in tasks]) is True
 
 
-def test_tune_tasks(onnx_mnist, tmpdir_factory):
-    pytest.importorskip("onnx")
+def test_tune_tasks(keras_simple, tmpdir_factory):
+    pytest.importorskip("tensorflow")
 
     tmpdir_name = tmpdir_factory.mktemp("data")
-    _autoscheduler_test_helper(onnx_mnist, tmpdir_name)
+    _autoscheduler_test_helper(keras_simple, tmpdir_name)
 
 
-def test_tune_tasks__tuning_records(onnx_mnist, tmpdir_factory):
-    pytest.importorskip("onnx")
+def test_tune_tasks__tuning_records(keras_simple, tmpdir_factory):
+    pytest.importorskip("tensorflow")
 
     tmpdir_name = tmpdir_factory.mktemp("data")
-    output_log_phase_1 = _autoscheduler_test_helper(onnx_mnist, tmpdir_name)
+    output_log_phase_1 = _autoscheduler_test_helper(keras_simple, tmpdir_name)
 
     # Exercises transfer learning by making sure a previous log exists
-    _autoscheduler_test_helper(onnx_mnist, tmpdir_name, prior_records=output_log_phase_1)
+    _autoscheduler_test_helper(keras_simple, tmpdir_name, prior_records=output_log_phase_1)
 
 
-def test_tune_tasks__no_early_stopping(onnx_mnist, tmpdir_factory):
-    pytest.importorskip("onnx")
+def test_tune_tasks__no_early_stopping(keras_simple, tmpdir_factory):
+    pytest.importorskip("tensorflow")
 
     tmpdir_name = tmpdir_factory.mktemp("data")
-    _autoscheduler_test_helper(onnx_mnist, tmpdir_name, early_stopping=None)
+    _autoscheduler_test_helper(keras_simple, tmpdir_name, early_stopping=None)

tests/python/driver/tvmc/test_command_line.py

@@ -20,26 +20,26 @@
 from tvm.driver.tvmc.main import _main
 
 
-def test_tvmc_cl_workflow(onnx_mnist, tmpdir_factory):
-    pytest.importorskip("onnx")
+def test_tvmc_cl_workflow(keras_simple, tmpdir_factory):
+    pytest.importorskip("tensorflow")
 
     tmpdir = tmpdir_factory.mktemp("data")
 
     # Test model tuning
-    log_path = os.path.join(tmpdir, "mnist-autotuner_records.json")
+    log_path = os.path.join(tmpdir, "keras-autotuner_records.json")
     tuning_str = (
         f"tvmc tune --target llvm --output {log_path} "
-        f"--trials 5 --enable-autoscheduler {onnx_mnist}"
+        f"--trials 2 --enable-autoscheduler {keras_simple}"
     )
     tuning_args = tuning_str.split(" ")[1:]
     _main(tuning_args)
     assert os.path.exists(log_path)
 
     # Test model compilation
-    package_path = os.path.join(tmpdir, "mnist-tvm.tar")
+    package_path = os.path.join(tmpdir, "keras-tvm.tar")
     compile_str = (
         f"tvmc compile --target llvm --tuning-records {log_path} "
-        f"--output {package_path} {onnx_mnist}"
+        f"--output {package_path} {keras_simple}"
     )
     compile_args = compile_str.split(" ")[1:]
     _main(compile_args)

tests/python/driver/tvmc/test_model.py

@@ -23,11 +23,11 @@
 from tvm.driver.tvmc.model import TVMCModel, TVMCPackage, TVMCResult
 
 
-def test_tvmc_workflow(onnx_mnist):
-    pytest.importorskip("onnx")
+def test_tvmc_workflow(keras_simple):
+    pytest.importorskip("tensorflow")
 
-    tvmc_model = tvmc.load(onnx_mnist)
-    tuning_records = tvmc.tune(tvmc_model, target="llvm", enable_autoscheduler=True, trials=5)
+    tvmc_model = tvmc.load(keras_simple)
+    tuning_records = tvmc.tune(tvmc_model, target="llvm", enable_autoscheduler=True, trials=2)
     tvmc_package = tvmc.compile(tvmc_model, tuning_records=tuning_records, target="llvm")
     result = tvmc.run(tvmc_package, device="cpu")
     assert type(tvmc_model) is TVMCModel
@@ -39,14 +39,14 @@ def test_tvmc_workflow(onnx_mnist):
     assert "output_0" in result.outputs.keys()
 
 
-def test_save_load_model(onnx_mnist, tmpdir_factory):
+def test_save_load_model(keras_simple, tmpdir_factory):
     pytest.importorskip("onnx")
 
     tmpdir = tmpdir_factory.mktemp("data")
-    tvmc_model = tvmc.load(onnx_mnist)
+    tvmc_model = tvmc.load(keras_simple)
 
     # Create tuning artifacts
-    tvmc.tune(tvmc_model, target="llvm", trials=4)
+    tvmc.tune(tvmc_model, target="llvm", trials=2)
 
     # Create package artifacts
     tvmc.compile(tvmc_model, target="llvm")