From 86126cdd60879a988fda3b729b8dffadcf5a3165 Mon Sep 17 00:00:00 2001 From: Gustavo Romero Date: Tue, 25 May 2021 10:13:38 -0300 Subject: [PATCH] [TVMC] Add support for the MLF to 'compile' command (#8086) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TVMC] Add support for the MLF to 'compile' command Add support for the Model Library Format (MLF) to 'tvmc' so users can output compilation artifacts to a MLF archive passing the new flag '--output-format mlf'. For instance: $ python3 -m tvm.driver.tvmc compile ./sine_model.tflite --target="c" --output sine.tar --output-format mlf will generate a sine.tar archive that is serialized accordingly to the MLF. Since the MLF is currently meant to be used only on micro targets, an error is generated if one tries to run a MLF outside a micro context. The micro context does not exist yet but will be later introduced as part of the [RFC] "TVMC: Add support for µTVM". That commit also adds 3 pytest tests to test tvmc + MLF. Finally, it also fixes some missing periods in the 'compile' command help sections and renames export_format to output_format so there is no confusion with flag '--dump-code', which contains "formats to export" in its help section. Signed-off-by: Gustavo Romero * Fix missing importorskip in the import_package test Fix missing importorskip() in the import_package test allowing the test in question to be skipped when 'tflite' is not installed in the test environment, otherwise the test will fail with: [...] > archive_path = exported_tvmc_package.package_path E AttributeError: 'str' object has no attribute 'package_path' --- python/tvm/driver/tvmc/compiler.py | 43 ++++++---- python/tvm/driver/tvmc/model.py | 93 ++++++++++++++++++---- python/tvm/driver/tvmc/runner.py | 8 ++ python/tvm/micro/model_library_format.py | 7 ++ tests/python/driver/tvmc/conftest.py | 25 +++++- tests/python/driver/tvmc/test_mlf.py | 99 ++++++++++++++++++++++++ 6 files changed, 241 insertions(+), 34 deletions(-) create mode 100644 tests/python/driver/tvmc/test_mlf.py 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."