Test for TVMC compile with TensorRT integration

apache · May 25, 2021 · e5b1a92 · e5b1a92
1 parent 20a86ca
commit e5b1a92
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/tests/python/driver/tvmc/test_compiler.py b/tests/python/driver/tvmc/test_compiler.py
@@ -330,6 +330,20 @@ def test_compile_tflite_module_with_external_codegen_vitis_ai(tflite_mobilenet_v
     assert type(tvmc_package.params) is bytearray
     assert os.path.exists(dumps_path)
 
+def test_compile_tflite_module_with_external_codegen_tensorrt(tflite_mobilenet_v1_0_25_128):
+    pytest.importorskip("tflite")
+
+    tvmc_model = tvmc.load(tflite_mobilenet_v1_0_25_128)
+    tvmc_package = tvmc.compiler.compile_model(
+        tvmc_model, target="tensorrt, llvm", dump_code="relay")
+    dumps_path = tvmc_package.package_path + ".relay"
+
+    # check for output types
+    assert type(tvmc_package) is TVMCPackage
+    assert type(tvmc_package.graph) is str
+    assert type(tvmc_package.lib_path) is str
+    assert type(tvmc_package.params) is bytearray
+    assert os.path.exists(dumps_path)
 
 @mock.patch("tvm.relay.build")
 @mock.patch("tvm.driver.tvmc.composite_target.get_codegen_by_target")

diff --git a/tests/python/driver/tvmc/test_composite_target.py b/tests/python/driver/tvmc/test_composite_target.py
@@ -35,6 +35,7 @@ def test_get_codegen_names():
 
     assert "ethos-n77" in names
     assert "vitis-ai" in names
+    assert "tensorrt" in names
     assert len(names) > 0