diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py index dcb770b9a563..0cc26872cd47 100644 --- a/python/tvm/driver/tvmc/compiler.py +++ b/python/tvm/driver/tvmc/compiler.py @@ -40,40 +40,48 @@ def add_compile_parser(subparsers): """ Include parser for 'compile' subcommand """ - parser = subparsers.add_parser("compile", help="compile a model") + parser = subparsers.add_parser("compile", help="compile a model.") parser.set_defaults(func=drive_compile) parser.add_argument( "--cross-compiler", default="", - help="the cross compiler to generate target libraries, e.g. 'aarch64-linux-gnu-gcc'", + help="the cross compiler to generate target libraries, e.g. 'aarch64-linux-gnu-gcc'.", ) parser.add_argument( "--cross-compiler-options", default="", - help="the cross compiler options to generate target libraries, e.g. '-mfpu=neon-vfpv4'", + help="the cross compiler options to generate target libraries, e.g. '-mfpu=neon-vfpv4'.", ) parser.add_argument( "--desired-layout", choices=["NCHW", "NHWC"], default=None, - help="change the data layout of the whole graph", + help="change the data layout of the whole graph.", ) parser.add_argument( "--dump-code", metavar="FORMAT", default="", - help="comma separarated list of formats to export, e.g. 'asm,ll,relay' ", + help="comma separated list of formats to export the input model, e.g. 'asm,ll,relay'.", ) parser.add_argument( "--model-format", choices=frontends.get_frontend_names(), - help="specify input model format", + help="specify input model format.", ) parser.add_argument( "-o", "--output", default="module.tar", - help="output the compiled module to an archive", + help="output the compiled module to a specifed archive. Defaults to 'module.tar'.", + ) + parser.add_argument( + "-f", + "--output-format", + choices=["so", "mlf"], + default="so", + help="output format. Use 'so' for shared object or 'mlf' for Model Library Format " + "(only for µTVM targets). Defaults to 'so'.", ) parser.add_argument( "--target", @@ -85,23 +93,23 @@ def add_compile_parser(subparsers): metavar="PATH", default="", help="path to an auto-tuning log file by AutoTVM. If not presented, " - "the fallback/tophub configs will be used", + "the fallback/tophub configs will be used.", ) - parser.add_argument("-v", "--verbose", action="count", default=0, help="increase verbosity") + parser.add_argument("-v", "--verbose", action="count", default=0, help="increase verbosity.") # TODO (@leandron) This is a path to a physical file, but # can be improved in future to add integration with a modelzoo # or URL, for example. - parser.add_argument("FILE", help="path to the input model file") + parser.add_argument("FILE", help="path to the input model file.") parser.add_argument( "--input-shapes", help="specify non-generic shapes for model to run, format is " - '"input_name:[dim1,dim2,...,dimn] input_name2:[dim1,dim2]"', + '"input_name:[dim1,dim2,...,dimn] input_name2:[dim1,dim2]".', type=common.parse_shape_string, default=None, ) parser.add_argument( "--disabled-pass", - help="disable specific passes, comma-separated list of pass names", + help="disable specific passes, comma-separated list of pass names.", type=common.parse_pass_list_str, default="", ) @@ -132,6 +140,7 @@ def drive_compile(args): package_path=args.output, cross=args.cross_compiler, cross_options=args.cross_compiler_options, + output_format=args.output_format, dump_code=dump_code, target_host=None, desired_layout=args.desired_layout, @@ -148,7 +157,7 @@ def compile_model( package_path: Optional[str] = None, cross: Optional[Union[str, Callable]] = None, cross_options: Optional[str] = None, - export_format: str = "so", + output_format: str = "so", dump_code: Optional[List[str]] = None, target_host: Optional[str] = None, desired_layout: Optional[str] = None, @@ -177,7 +186,7 @@ def compile_model( Function that performs the actual compilation cross_options : str, optional Command line options to be passed to the cross compiler. - export_format : str + output_format : str What format to use when saving the function library. Must be one of "so" or "tar". When compiling for a remote device without a cross compiler, "tar" will likely work better. dump_code : list, optional @@ -262,7 +271,11 @@ def compile_model( # Create a new tvmc model package object from the graph definition. package_path = tvmc_model.export_package( - graph_module, package_path, cross, cross_options, export_format + graph_module, + package_path, + cross, + cross_options, + output_format, ) # Write dumps to file. diff --git a/python/tvm/driver/tvmc/model.py b/python/tvm/driver/tvmc/model.py index 26a1e3600b96..d9e3266c766f 100644 --- a/python/tvm/driver/tvmc/model.py +++ b/python/tvm/driver/tvmc/model.py @@ -53,6 +53,7 @@ from tvm import relay from tvm.contrib import utils from tvm.relay.backend.executor_factory import GraphExecutorFactoryModule +from tvm.micro import export_model_library_format from .common import TVMCException @@ -175,7 +176,7 @@ def default_package_path(self): """ return self._tmp_dir.relpath("model_package.tar") - def export_package( + def export_classic_format( self, executor_factory: GraphExecutorFactoryModule, package_path: Optional[str] = None, @@ -203,8 +204,6 @@ def export_package( package_path : str The path that the package was saved to. """ - if lib_format not in ["so", "tar"]: - raise TVMCException("Only .so and .tar export formats are supported.") lib_name = "mod." + lib_format graph_name = "mod.json" param_name = "mod.params" @@ -241,6 +240,50 @@ def export_package( return package_path + def export_package( + self, + executor_factory: GraphExecutorFactoryModule, + package_path: Optional[str] = None, + cross: Optional[Union[str, Callable]] = None, + cross_options: Optional[str] = None, + output_format: str = "so", + ): + """Save this TVMCModel to file. + Parameters + ---------- + executor_factory : GraphExecutorFactoryModule + The factory containing compiled the compiled artifacts needed to run this model. + package_path : str, None + Where the model should be saved. Note that it will be packaged as a .tar file. + If not provided, the package will be saved to a generically named file in tmp. + cross : str or callable object, optional + Function that performs the actual compilation. + cross_options : str, optional + Command line options to be passed to the cross compiler. + output_format : str + How to save the modules function library. Must be one of "so" and "tar" to save + using the classic format or "mlf" to save using the Model Library Format. + + Returns + ------- + package_path : str + The path that the package was saved to. + """ + if output_format not in ["so", "tar", "mlf"]: + raise TVMCException("Only 'so', 'tar', and 'mlf' output formats are supported.") + + if output_format == "mlf" and cross: + raise TVMCException("Specifying the MLF output and a cross compiler is not supported.") + + if output_format in ["so", "tar"]: + package_path = self.export_classic_format( + executor_factory, package_path, cross, cross_options, output_format + ) + elif output_format == "mlf": + package_path = export_model_library_format(executor_factory, package_path) + + return package_path + def summary(self, file: TextIO = None): """Print the IR corressponding to this model. @@ -274,25 +317,41 @@ def import_package(self, package_path: str): package_path : str The path to the saved TVMCPackage. """ - lib_name_so = "mod.so" - lib_name_tar = "mod.tar" - graph_name = "mod.json" - param_name = "mod.params" - temp = self._tmp_dir t = tarfile.open(package_path) t.extractall(temp.relpath(".")) - with open(temp.relpath(param_name), "rb") as param_file: - self.params = bytearray(param_file.read()) - self.graph = open(temp.relpath(graph_name)).read() - if os.path.exists(temp.relpath(lib_name_so)): - self.lib_name = lib_name_so - elif os.path.exists(temp.relpath(lib_name_tar)): - self.lib_name = lib_name_tar + if os.path.exists(temp.relpath("metadata.json")): + # Model Library Format (MLF) + self.lib_name = None + self.lib_path = None + + graph = temp.relpath("runtime-config/graph/graph.json") + params = temp.relpath("parameters/default.params") + + self.type = "mlf" else: - raise TVMCException("Couldn't find exported library in the package.") - self.lib_path = temp.relpath(self.lib_name) + # Classic format + lib_name_so = "mod.so" + lib_name_tar = "mod.tar" + if os.path.exists(temp.relpath(lib_name_so)): + self.lib_name = lib_name_so + elif os.path.exists(temp.relpath(lib_name_tar)): + self.lib_name = lib_name_tar + else: + raise TVMCException("Couldn't find exported library in the package.") + self.lib_path = temp.relpath(self.lib_name) + + graph = temp.relpath("mod.json") + params = temp.relpath("mod.params") + + self.type = "classic" + + with open(params, "rb") as param_file: + self.params = bytearray(param_file.read()) + + with open(graph) as graph_file: + self.graph = graph_file.read() class TVMCResult(object): diff --git a/python/tvm/driver/tvmc/runner.py b/python/tvm/driver/tvmc/runner.py index c59689face63..ba0f7d2c2d6c 100644 --- a/python/tvm/driver/tvmc/runner.py +++ b/python/tvm/driver/tvmc/runner.py @@ -359,6 +359,14 @@ def run_module( "Try calling tvmc.compile on the model before running it." ) + # Currently only two package formats are supported: "classic" and + # "mlf". The later can only be used for micro targets, i.e. with µTVM. + if tvmc_package.type == "mlf": + raise TVMCException( + "You're trying to run a model saved using the Model Library Format (MLF)." + "MLF can only be used to run micro targets (µTVM)." + ) + if hostname: if isinstance(port, str): port = int(port) diff --git a/python/tvm/micro/model_library_format.py b/python/tvm/micro/model_library_format.py index be991e22a0f8..1cc3adf9ae07 100644 --- a/python/tvm/micro/model_library_format.py +++ b/python/tvm/micro/model_library_format.py @@ -216,6 +216,11 @@ def export_model_library_format(mod: executor_factory.ExecutorFactoryModule, fil The return value of tvm.relay.build, which will be exported into Model Library Format. file_name : str Path to the .tar archive to generate. + + Returns + ------- + file_name : str + The path to the generated .tar archive. """ tempdir = utils.tempdir() is_aot = isinstance(mod, executor_factory.AOTExecutorFactoryModule) @@ -260,3 +265,5 @@ def reset(tarinfo): return tarinfo tar_f.add(tempdir.temp_dir, arcname=".", filter=reset) + + return file_name diff --git a/tests/python/driver/tvmc/conftest.py b/tests/python/driver/tvmc/conftest.py index f7cbf92bca30..9c0d8fa8911e 100644 --- a/tests/python/driver/tvmc/conftest.py +++ b/tests/python/driver/tvmc/conftest.py @@ -41,7 +41,7 @@ def download_and_untar(model_url, model_sub_path, temp_dir): return os.path.join(temp_dir, model_sub_path) -def get_sample_compiled_module(target_dir, package_filename): +def get_sample_compiled_module(target_dir, package_filename, output_format="so"): """Support function that returns a TFLite compiled module""" base_url = "https://storage.googleapis.com/download.tensorflow.org/models" model_url = "mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz" @@ -53,7 +53,10 @@ def get_sample_compiled_module(target_dir, package_filename): tvmc_model = tvmc.frontends.load_model(model_file) return tvmc.compiler.compile_model( - tvmc_model, target="llvm", package_path=os.path.join(target_dir, package_filename) + tvmc_model, + target="llvm", + package_path=os.path.join(target_dir, package_filename), + output_format=output_format, ) @@ -182,6 +185,24 @@ def tflite_compiled_model(tmpdir_factory): return get_sample_compiled_module(target_dir, "mock.tar") +@pytest.fixture(scope="session") +def tflite_compiled_model_mlf(tmpdir_factory): + + # Not all CI environments will have TFLite installed + # so we need to safely skip this fixture that will + # crash the tests that rely on it. + # As this is a pytest.fixture, we cannot take advantage + # of pytest.importorskip. Using the block below instead. + try: + import tflite + except ImportError: + print("Cannot import tflite, which is required by tflite_compiled_module_as_tarfile.") + return "" + + target_dir = tmpdir_factory.mktemp("data") + return get_sample_compiled_module(target_dir, "mock.tar", "mlf") + + @pytest.fixture(scope="session") def imagenet_cat(tmpdir_factory): tmpdir_name = tmpdir_factory.mktemp("data") diff --git a/tests/python/driver/tvmc/test_mlf.py b/tests/python/driver/tvmc/test_mlf.py new file mode 100644 index 000000000000..48be5a810bc5 --- /dev/null +++ b/tests/python/driver/tvmc/test_mlf.py @@ -0,0 +1,99 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pytest +import os + +import tvm +from tvm.driver import tvmc +from tvm.driver.tvmc.main import _main +from tvm.driver.tvmc.model import TVMCPackage, TVMCException + + +def test_tvmc_cl_compile_run_mlf(tflite_mobilenet_v1_1_quant, tmpdir_factory): + pytest.importorskip("tflite") + + output_dir = tmpdir_factory.mktemp("mlf") + input_model = tflite_mobilenet_v1_1_quant + output_file = os.path.join(output_dir, "mock.tar") + + # Compile the input model and generate a Model Library Format (MLF) archive. + tvmc_cmd = ( + f"tvmc compile {input_model} --target='llvm' --output {output_file} --output-format mlf" + ) + tvmc_args = tvmc_cmd.split(" ")[1:] + _main(tvmc_args) + assert os.path.exists(output_file), "Could not find the exported MLF archive." + + # Run the MLF archive. It must fail since it's only supported on micro targets. + tvmc_cmd = f"tvmc run {output_file}" + tvmc_args = tvmc_cmd.split(" ")[1:] + exit_code = _main(tvmc_args) + on_error = "Trying to run a MLF archive must fail because it's only supported on micro targets." + assert exit_code != 0, on_error + + +def test_tvmc_export_package_mlf(tflite_mobilenet_v1_1_quant, tmpdir_factory): + pytest.importorskip("tflite") + + tvmc_model = tvmc.frontends.load_model(tflite_mobilenet_v1_1_quant) + mod, params = tvmc_model.mod, tvmc_model.params + + graph_module = tvm.relay.build(mod, target="llvm", params=params) + + output_dir = tmpdir_factory.mktemp("mlf") + output_file = os.path.join(output_dir, "mock.tar") + + # Try to export MLF with no cross compiler set. No exception must be thrown. + tvmc_model.export_package( + executor_factory=graph_module, + package_path=output_file, + cross=None, + output_format="mlf", + ) + assert os.path.exists(output_file), "Could not find the exported MLF archive." + + # Try to export a MLF whilst also specifying a cross compiler. Since + # that's not supported it must throw a TVMCException and report the + # reason accordingly. + with pytest.raises(TVMCException) as exp: + tvmc_model.export_package( + executor_factory=graph_module, + package_path=output_file, + cross="cc", + output_format="mlf", + ) + expected_reason = "Specifying the MLF output and a cross compiler is not supported." + on_error = "A TVMCException was caught but its reason is not the expected one." + assert str(exp.value) == expected_reason, on_error + + +def test_tvmc_import_package_mlf(tflite_compiled_model_mlf): + pytest.importorskip("tflite") + + # Compile and export a model to a MLF archive so it can be imported. + exported_tvmc_package = tflite_compiled_model_mlf + archive_path = exported_tvmc_package.package_path + + # Import the MLF archive. TVMCPackage constructor will call import_package method. + tvmc_package = TVMCPackage(archive_path) + + assert tvmc_package.lib_name is None, ".lib_name must not be set in the MLF archive." + assert tvmc_package.lib_path is None, ".lib_path must not be set in the MLF archive." + assert tvmc_package.graph is not None, ".graph must be set in the MLF archive." + assert tvmc_package.params is not None, ".params must be set in the MLF archive." + assert tvmc_package.type == "mlf", ".type must be set to 'mlf' in the MLF format."