areusch commented on a change in pull request #7823:
URL: https://github.com/apache/tvm/pull/7823#discussion_r616771199
##########
File path: python/tvm/driver/tvmc/autotuner.py
##########
@@ -255,97 +380,100 @@ def drive_tune(args):
# min_repeat_ms should be:
# a. the value provided by the user, if any, or
# b. 0ms in case target is "cpu"; otherwise 1000ms
- if args.min_repeat_ms is not None:
- min_repeat_ms = args.min_repeat_ms
- else:
+ if min_repeat_ms is None:
min_repeat_ms = 0 if target.keys[0] == "cpu" else 1000
logger.debug("Default --min-repeat-ms for this target is %s",
min_repeat_ms)
- if args.rpc_tracker:
- runner_ctor = auto_scheduler.RPCRunner if args.enable_autoscheduler
else autotvm.RPCRunner
+ if rpc_tracker:
+ runner_ctor = auto_scheduler.RPCRunner if enable_autoscheduler else
autotvm.RPCRunner
runner = runner_ctor(
- key=args.rpc_key,
+ key=rpc_key,
host=rpc_hostname,
port=rpc_port,
- number=args.number,
- repeat=args.repeat,
- n_parallel=args.parallel,
- timeout=args.timeout,
+ number=number,
+ repeat=repeat,
+ n_parallel=parallel,
+ timeout=timeout,
min_repeat_ms=min_repeat_ms,
)
else:
logger.info("starting localhost tuning")
runner_ctor = (
- auto_scheduler.LocalRunner if args.enable_autoscheduler else
autotvm.LocalRunner
+ auto_scheduler.LocalRPCMeasureContext if enable_autoscheduler else
autotvm.LocalRunner
)
runner = runner_ctor(
- number=args.number,
- repeat=args.repeat,
- timeout=args.timeout,
+ number=number,
+ repeat=repeat,
+ timeout=timeout,
min_repeat_ms=min_repeat_ms,
)
- if args.enable_autoscheduler:
- # Specify hardware parameters
- hardware_params = auto_scheduler.HardwareParams(
- args.num_cores,
- args.vector_unit_bytes,
- args.cache_line_bytes,
- args.max_shared_memory_per_block,
- args.max_local_memory_per_block,
- args.max_threads_per_block,
- args.max_vthread_extent,
- args.warp_size,
- )
+ if enable_autoscheduler:
+
tasks, weights = autoscheduler_get_tuning_tasks(
mod=mod,
params=params,
target=target,
- alter_layout=args.desired_layout,
+ alter_layout=desired_layout,
hardware_params=hardware_params,
- include_simple_tasks=args.include_simple_tasks,
+ include_simple_tasks=include_simple_tasks,
)
+ # If not specified, choose a number of trials likely to produce good
results.
+ if trials is None:
+ trials = 10000
Review comment:
i'm ok with it being a default. I also like the table-of-config-values
idea.
##########
File path: python/tvm/driver/tvmc/autotuner.py
##########
@@ -255,97 +388,113 @@ def drive_tune(args):
# min_repeat_ms should be:
# a. the value provided by the user, if any, or
# b. 0ms in case target is "cpu"; otherwise 1000ms
- if args.min_repeat_ms is not None:
- min_repeat_ms = args.min_repeat_ms
- else:
+ if min_repeat_ms is None:
min_repeat_ms = 0 if target.keys[0] == "cpu" else 1000
logger.debug("Default --min-repeat-ms for this target is %s",
min_repeat_ms)
- if args.rpc_tracker:
- runner_ctor = auto_scheduler.RPCRunner if args.enable_autoscheduler
else autotvm.RPCRunner
+ if rpc_key:
+ if hostname is None or port is None:
+ raise common.TVMCException(
+ "You must provide a hostname and port to connect to a remote
RPC device."
+ )
+ if isinstance(port, str):
+ port = int(port)
+
+ runner_ctor = auto_scheduler.RPCRunner if enable_autoscheduler else
autotvm.RPCRunner
runner = runner_ctor(
- key=args.rpc_key,
- host=rpc_hostname,
- port=rpc_port,
- number=args.number,
- repeat=args.repeat,
- n_parallel=args.parallel,
- timeout=args.timeout,
+ key=rpc_key,
+ host=hostname,
+ port=port,
+ number=number,
+ repeat=repeat,
+ n_parallel=parallel,
+ timeout=timeout,
min_repeat_ms=min_repeat_ms,
)
else:
logger.info("starting localhost tuning")
runner_ctor = (
- auto_scheduler.LocalRunner if args.enable_autoscheduler else
autotvm.LocalRunner
+ auto_scheduler.LocalRPCMeasureContext if enable_autoscheduler else
autotvm.LocalRunner
)
- runner = runner_ctor(
- number=args.number,
- repeat=args.repeat,
- timeout=args.timeout,
+ local_server = runner_ctor(
+ number=number,
+ repeat=repeat,
+ timeout=timeout,
min_repeat_ms=min_repeat_ms,
)
- if args.enable_autoscheduler:
- # Specify hardware parameters
- hardware_params = auto_scheduler.HardwareParams(
- args.num_cores,
- args.vector_unit_bytes,
- args.cache_line_bytes,
- args.max_shared_memory_per_block,
- args.max_local_memory_per_block,
- args.max_threads_per_block,
- args.max_vthread_extent,
- args.warp_size,
- )
+ # For autoscheduling on some devices, we need to maintain a
LocalRPCMeasureContext object.
+ if enable_autoscheduler:
+ runner = local_server.runner
+ else:
+ runner = local_server
+
+ if enable_autoscheduler:
+
tasks, weights = autoscheduler_get_tuning_tasks(
mod=mod,
params=params,
target=target,
- alter_layout=args.desired_layout,
+ alter_layout=desired_layout,
hardware_params=hardware_params,
- include_simple_tasks=args.include_simple_tasks,
+ include_simple_tasks=include_simple_tasks,
)
+ # If not specified, choose a number of trials likely to produce good
results.
+ if trials is None:
+ trials = 10000
+
# Create the autoscheduler tuning options
tuning_options = auto_scheduler.TuningOptions(
- num_measure_trials=args.trials,
- measure_callbacks=[auto_scheduler.RecordToFile(args.output)],
+ num_measure_trials=trials,
+ measure_callbacks=[auto_scheduler.RecordToFile(tuning_records)],
runner=runner,
- early_stopping=args.early_stopping,
+ early_stopping=early_stopping,
)
# Schedule the tasks (i.e., produce a schedule for each task)
- schedule_tasks(
- tasks, weights, tuning_options, args.tuning_records,
args.log_estimated_latency
- )
+ schedule_tasks(tasks, weights, tuning_options, prior_records,
log_estimated_latency)
else:
tasks = autotvm_get_tuning_tasks(
mod=mod,
params=params,
target=target,
- alter_layout=args.desired_layout,
+ alter_layout=desired_layout,
)
+ # If trails isn't specified, default to a number likely to produce good
+ # results without taking too much time.
+ if trials is None:
+ trials = 1000
Review comment:
I think it depends whether the user understands what task extraction
does. if not, I think trials-per-task is intuitive. I understand that total
trials better implies the runtime, but it's hard to estimate that if you don't
know how task extraction behaves.
perhaps we need a tvmc subcommand which just runs and prints task extraction.
##########
File path: python/tvm/driver/tvmc/model.py
##########
@@ -0,0 +1,372 @@
+# 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.
+"""
+This file contains the definition of a set of classes that wrap the outputs
+of TVMC functions to create a simpler and more intuitive API.
+
+There is one class for each required stage of a TVM workflow.
+The TVMCModel represents the result of importing a model into TVM, it
+contains the precompiled graph definition and parameters that define
+what the model does.
+
+Compiling a TVMCModel produces a TVMCPackage, which contains the generated
+artifacts that allow the model to be run on the target hardware.
+
+Running a TVMCPackage produces a TVMCResult, which contains the outputs of
+the model and the measured runtime.
+
+Examples
+--------
+The following code shows a full lifecycle for a model using tvmc, first the
+model is imported from an exterior framework, in this case onnx, then it
+is tuned to find the best schedules on CPU, then compiled into a TVMCPackage,
+and finally run.
+
+.. code-block:: python
+ tvmc_model = tvmc.load("my_model.onnx")
+ tuning_records = tvmc.tune(tvmc_model, target="llvm")
+ tvmc_package = tvmc.compile(tvmc_model, target="llvm",
tuning_records=tuning_records)
+ result = tvmc.run(tvmc_package, device="cpu")
+ print(result)
+"""
+import os
+import tarfile
+from typing import Optional, Union, List, Dict, Callable, TextIO
+import numpy as np
+
+import tvm
+import tvm.contrib.cc
+from tvm import relay
+from tvm.contrib import utils
+from tvm.relay.backend.graph_executor_factory import GraphExecutorFactoryModule
+
+from .common import TVMCException
+
+
+class TVMCModel(object):
+ """Initialize a TVMC model from a relay model definition or a saved file.
+
+ Parameters
+ ----------
+ mod : tvm.IRModule, optional
+ The relay module corresponding to this model.
+ params : dict, optional
+ A parameter dictionary for the model.
+ model_path: str, optional
+ An alternative way to load a TVMCModel, the path to a previously
+ saved model.
+ """
+
+ def __init__(
+ self,
+ mod: Optional[tvm.IRModule] = None,
+ params: Optional[Dict[str, tvm.nd.NDArray]] = None,
+ model_path: Optional[str] = None,
+ ):
+ if (mod is None or params is None) and (model_path is None):
+ raise TVMCException(
+ "Either mod and params must be provided "
+ "or a path to a previously saved TVMCModel"
+ )
+ self._tmp_dir = utils.tempdir()
+ if model_path is not None:
+ self.load(model_path)
+ else:
+ self.mod = mod
+ self.params = params if params else {}
+
+ def save(self, model_path: str):
+ """Save the TVMCModel to disk.
+
+ Note that this saves the graph representation,
+ the parameters, and the tuning records if applicable. It will not save
any
+ compiled artifacts.
+
+ Parameters
+ ----------
+ model_path : str
+ A full path to save this TVMCModel to including the output file
name.
+ The file will be saved as a tar file so using a ".tar" extension
is advised.
+ """
+ temp = self._tmp_dir
+
+ # Save relay graph
+ relay_name = "model.json"
+ relay_path = temp.relpath(relay_name)
+ with open(relay_path, "w") as relay_file:
+ relay_file.write(tvm.ir.save_json(self.mod))
+
+ # Save params
+ params_name = "model.params"
+ params_path = temp.relpath(params_name)
+ with open(params_path, "wb") as params_file:
+ params_file.write(relay.save_param_dict(self.params))
+
+ # Create a tar file.
+ with tarfile.open(model_path, "w") as tar:
+ tar.add(relay_path, relay_name)
+ tar.add(params_path, params_name)
+
+ def load(self, model_path: str):
+ """Load a TVMCModel from disk.
+
+ Parameters
+ ----------
+ model_path : str
+ A path to load the TVMCModel from.
+ """
+ temp = self._tmp_dir
+ t = tarfile.open(model_path)
+ t.extractall(temp.relpath("."))
+
+ # Load relay IR.
+ relay_path = temp.relpath("model.json")
+ with open(relay_path, "r") as relay_file:
+ self.mod = tvm.ir.load_json(relay_file.read())
+
+ # Load parameter dictionary.
+ params_path = temp.relpath("model.params")
+ with open(params_path, "rb") as params_file:
+ self.params = relay.load_param_dict(params_file.read())
+
+ def get_temp_path(self, file_name: str):
+ """Get the full path for a filename stored in this model's temp
directory.
+
+ Parameters
+ ----------
+ file_name : str
+ The name of the file within this model's temp directory.
+
+ Returns
+ -------
+ temp_path : str
+ A path to a file in this model's temporary directory.
+ """
+ return self._tmp_dir.relpath(file_name)
+
+ def export_package(
+ self,
+ executor_factory: GraphExecutorFactoryModule,
+ package_path: Optional[str] = None,
+ cross: Optional[Union[str, Callable]] = None,
+ lib_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.
+ lib_format : str
+ How to export the modules function library. Must be one of "so" or
"tar".
+
+ Returns
+ -------
+ 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"
+
+ temp = self._tmp_dir
+ if package_path is None:
+ package_path = temp.relpath("model_package.tar")
+ path_lib = temp.relpath(lib_name)
+
+ if not cross:
+ executor_factory.get_lib().export_library(path_lib)
+ else:
+ executor_factory.get_lib().export_library(
+ path_lib, tvm.contrib.cc.cross_compiler(cross)
+ )
+ self.lib_path = path_lib
+
+ with open(temp.relpath(graph_name), "w") as graph_file:
+ graph_file.write(executor_factory.get_json())
+
+ with open(temp.relpath(param_name), "wb") as params_file:
+
params_file.write(relay.save_param_dict(executor_factory.get_params()))
+
+ # Package up all the temp files into a tar file.
+ with tarfile.open(package_path, "w") as tar:
+ tar.add(path_lib, lib_name)
+ tar.add(temp.relpath(graph_name), graph_name)
+ tar.add(temp.relpath(param_name), param_name)
+
+ return package_path
+
+ def summary(self, file: TextIO = None):
Review comment:
I guess this is meant to correspond to tensorflow summary()?
##########
File path: python/tvm/driver/tvmc/autotuner.py
##########
@@ -228,24 +240,137 @@ def drive_tune(args):
args: argparse.Namespace
Arguments from command line parser.
"""
- # extra arguments validation before importing the model, so that obvious
errors
- # are pointed in advance.
- if args.rpc_tracker:
- parsed_url = urlparse("//%s" % args.rpc_tracker)
+ tvmc_model = frontends.load_model(args.FILE, args.model_format,
shape_dict=args.input_shapes)
+ tvmc_model.tuning_records = args.tuning_records
+ # Specify hardware parameters, although they'll only be used if
autoscheduling.
+ hardware_params = auto_scheduler.HardwareParams(
+ args.num_cores,
+ args.vector_unit_bytes,
+ args.cache_line_bytes,
+ args.max_shared_memory_per_block,
+ args.max_local_memory_per_block,
+ args.max_threads_per_block,
+ args.max_vthread_extent,
+ args.warp_size,
+ args.target,
+ args.target_host,
+ )
+
+ tune_model(
+ tvmc_model,
+ args.target,
+ args.output,
+ args.enable_autoscheduler,
+ args.rpc_key,
+ args.rpc_tracker,
+ args.trials,
+ args.target_host,
+ args.tuner,
+ args.min_repeat_ms,
+ args.early_stopping,
+ args.desired_layout,
+ args.timeout,
+ args.number,
+ args.repeat,
+ args.parallel,
+ hardware_params,
+ args.include_simple_tasks,
+ args.log_estimated_latency,
+ )
+
+
+def tune_model(
+ tvmc_model: TVMCModel,
+ target: str,
+ tuning_records: Optional[str] = None,
+ enable_autoscheduler: bool = False,
+ rpc_key: Optional[str] = None,
+ rpc_tracker: Optional[str] = None,
+ trials: Optional[int] = None,
+ target_host: str = "llvm",
+ tuner: str = "xgb",
+ min_repeat_ms: Optional[int] = None,
+ early_stopping: Optional[int] = None,
+ desired_layout: Optional[str] = None,
+ timeout: int = 10,
+ number: int = 10,
+ repeat: int = 1,
+ parallel: int = 4,
+ hardware_params: Optional[HardwareParams] = None,
+ include_simple_tasks: bool = False,
+ log_estimated_latency: bool = False,
+):
+ """Use tuning to automatically optimize the functions in a model.
+
+ Parameters
+ ----------
+ tvmc_model : TVMCModel
+ The model to be optimized.
+ target : str
+ Compilation target as plain string, inline JSON or path to a JSON file.
+ tuning_records: str, optional
+ The path to a file that tuning results will be saved to. If not
specified,
+ a temporary file will be used.
Review comment:
I agree it's misleading to allow users to run tune without making sure
they've specified a place to put the result. I don't see how it would make
sense for the command-line interface to allow that.
one question if we unify though: i'd think if we do that, we should modify
some logic to place a delineating comment in the tuning log before we write the
new results.
##########
File path: python/tvm/driver/tvmc/autotuner.py
##########
@@ -228,24 +240,137 @@ def drive_tune(args):
args: argparse.Namespace
Arguments from command line parser.
"""
- # extra arguments validation before importing the model, so that obvious
errors
- # are pointed in advance.
- if args.rpc_tracker:
- parsed_url = urlparse("//%s" % args.rpc_tracker)
+ tvmc_model = frontends.load_model(args.FILE, args.model_format,
shape_dict=args.input_shapes)
+ tvmc_model.tuning_records = args.tuning_records
+ # Specify hardware parameters, although they'll only be used if
autoscheduling.
+ hardware_params = auto_scheduler.HardwareParams(
+ args.num_cores,
+ args.vector_unit_bytes,
+ args.cache_line_bytes,
+ args.max_shared_memory_per_block,
+ args.max_local_memory_per_block,
+ args.max_threads_per_block,
+ args.max_vthread_extent,
+ args.warp_size,
+ args.target,
+ args.target_host,
+ )
+
+ tune_model(
+ tvmc_model,
+ args.target,
+ args.output,
+ args.enable_autoscheduler,
+ args.rpc_key,
+ args.rpc_tracker,
+ args.trials,
+ args.target_host,
+ args.tuner,
+ args.min_repeat_ms,
+ args.early_stopping,
+ args.desired_layout,
+ args.timeout,
+ args.number,
+ args.repeat,
+ args.parallel,
+ hardware_params,
+ args.include_simple_tasks,
+ args.log_estimated_latency,
+ )
+
+
+def tune_model(
+ tvmc_model: TVMCModel,
+ target: str,
+ tuning_records: Optional[str] = None,
+ enable_autoscheduler: bool = False,
+ rpc_key: Optional[str] = None,
+ rpc_tracker: Optional[str] = None,
+ trials: Optional[int] = None,
+ target_host: str = "llvm",
+ tuner: str = "xgb",
+ min_repeat_ms: Optional[int] = None,
+ early_stopping: Optional[int] = None,
+ desired_layout: Optional[str] = None,
+ timeout: int = 10,
+ number: int = 10,
+ repeat: int = 1,
+ parallel: int = 4,
+ hardware_params: Optional[HardwareParams] = None,
+ include_simple_tasks: bool = False,
+ log_estimated_latency: bool = False,
+):
+ """Use tuning to automatically optimize the functions in a model.
+
+ Parameters
+ ----------
+ tvmc_model : TVMCModel
+ The model to be optimized.
+ target : str
+ Compilation target as plain string, inline JSON or path to a JSON file.
+ tuning_records: str, optional
+ The path to a file that tuning results will be saved to. If not
specified,
+ a temporary file will be used.
+ enable_autoscheduler : bool, optional
+ When true, use autoscheduling rather than autotvm. This should produce
+ faster kernels for compatible model-target pairs.
+ rpc_key : str, optional
+ The RPC tracker key of the target device. Required when rpc_tracker is
provided.
+ rpc_tracker : str, optional
+ The hostname and port (optional, defaults to 9090) of the RPC tracker,
+ e.g. 192.168.0.100:9999.
+ trials : int, optional
+ The number of schedules to try out. For autotvm, each task will have
this many
+ options explored. For autoscheduling, the total number of schedules
checked in
+ the entire model will be this many.
+ target_host : str, optional
+ The host compilation target, defaults to 'llvm'.
+ tuner : str, optional
+ The type of tuner to use when tuning with autotvm. Can be one of
+ "ga", "gridsearch", "random", "xgb", "xgb_knob", and "xgb-rank".
+ min_repeat_ms : int, optional
+ Minimum time to run each trial. Defaults to 0 on x86 and 1000 on other
targets.
+ early_stopping : int, optional
+ When specified, stop tuning after this number of trials if results
aren't improving.
+ desired_layout : str, optional
+ Can be one of "NCHW" or "NHWC". When specified, the graph will be
converted to this layout.
+ timeout : int, optional,
+ If a kernel trial lasts longer than this duration in seconds, it will
be
+ considered a failure.
+ number : int, optional
Review comment:
I guess to play devil's advocate: this PR is introducing a top-level
API. why is that out of scope now? I understand that it might be hard to do all
of that plus e.g. yaml considerations, but organization seems like a
fundamental part of this PR.
##########
File path: python/tvm/driver/tvmc/model.py
##########
@@ -0,0 +1,357 @@
+# 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.
+"""
+This file contains the definition of a set of classes that wrap the outputs
+of TVMC functions to create a simpler and more intuitive API.
+
+There is one class for each required stage of a TVM workflow.
+The TVMCModel represents the result of importing a model into TVM, it
+contains the precompiled graph definition and parameters that define
+what the model does.
+
+Compiling a TVMCModel produces a TVMCPackage, which contains the generated
+artifacts that allow the model to be run on the target hardware.
+
+Running a TVMCPackage produces a TVMCResult, which contains the outputs of
+the model and the measured runtime.
+
+Examples
+--------
+The following code shows a full lifecycle for a model using tvmc, first the
+model is imported from an exterior framework, in this case onnx, then it
+is tuned to find the best schedules on CPU, then compiled into a TVMCPackage,
+and finally run.
+
+.. code-block:: python
+ tvmc_model = tvmc.load("my_model.onnx")
+ tvmc_model = tvmc.tune(tvmc_model, target="llvm")
+ tvmc_package = tvmc.compile(tvmc_model, "llvm")
+ result = tvmc.run(tvmc_package, device="cpu")
+ print(result)
+"""
+import os
+import json
+import tarfile
+from typing import Optional, Union, List, Dict, Callable
+import numpy as np
+
+import tvm
+from tvm import relay
+from tvm.micro import export_model_library_format
+from tvm.contrib import utils
+from tvm.relay.backend.graph_executor_factory import GraphExecutorFactoryModule
+
+from .common import TVMCException
+
+
+class TVMCModel(object):
+ """Initialize a TVMC model from a relay model definition or a saved file.
+
+ Parameters
+ ----------
+ mod : tvm.IRModule, optional
+ The relay module corresponding to this model.
+ params : dict, optional
+ A parameter dictionary for the model.
+ model_path: str, optional
+ An alternative way to load a TVMCModel, the path to a previously
+ saved model.
+ name : str, optional
+ An optional name for the main library being compiled. If not specified,
+ 'default' will be used.
+ """
+
+ def __init__(
+ self,
+ mod: Optional[tvm.IRModule] = None,
+ params: Optional[Dict[str, tvm.nd.NDArray]] = None,
+ model_path: Optional[str] = None,
+ name: Optional[str] = None,
+ ):
+ if (mod is None or params is None) and (model_path is None):
+ raise TVMCException(
+ "Either mod and params must be provided "
+ "or a path to a previously saved TVMCModel"
+ )
+ self.mod = mod
+ self.params = params if params else {}
+ self.name = "default" if name is None else name
+ self.dumps = None
+ self.tuning_records = None
+ self.package_path = None
+ self._tmp_dir = utils.tempdir()
+ if model_path is not None:
+ self.load(model_path)
+
+ def save(self, model_path: str):
+ """Save the TVMCModel to disk. Note that this saves the graph
representation,
+ the parameters, and the tuning records if applicable. It will not save
any
+ compiled artifacts.
+
+ Parameters
+ ----------
+ model_path : str
+ A path to save this TVMCModel to.
+ """
+ temp = self._tmp_dir
+ metadata = {"model_name": self.name}
+
+ # Save metadata
+ metadata_name = "metadata.json"
+ metadata_path = temp.relpath(metadata_name)
+ with open(metadata_path, "w") as metadata_file:
+ json.dump(metadata, metadata_file, indent=2, sort_keys=True)
+
+ # Save relay graph
+ relay_name = "model.json"
+ relay_path = temp.relpath(relay_name)
+ with open(relay_path, "w") as relay_file:
+ relay_file.write(tvm.ir.save_json(self.mod))
+
+ # Save params
+ params_name = "model.params"
+ params_path = temp.relpath(params_name)
+ with open(params_path, "wb") as params_file:
+ params_file.write(relay.save_param_dict(self.params))
+
+ # Save tuning logs if they are being kept as part of this model.
+ records_name = "tuning_records"
+ records_path = self.get_default_tuning_path()
+
+ # Create a tar file.
+ with tarfile.open(model_path, "w") as tar:
+ tar.add(metadata_path, metadata_name)
+ tar.add(relay_path, relay_name)
+ tar.add(params_path, params_name)
+ if self.tuning_records == records_path:
+ tar.add(records_path, records_name)
+
+ def load(self, model_path: str):
+ """Load a TVMCModel from disk.
+
+ Parameters
+ ----------
+ model_path : str
+ A path to load the TVMCModel from.
+ """
+ temp = self._tmp_dir
+ t = tarfile.open(model_path)
+ t.extractall(temp.relpath("."))
+
+ # Load metadata.
+ metadata_path = temp.relpath("metadata.json")
+ with open(metadata_path, "r") as metadata_file:
+ metadata = json.load(metadata_file)
+ self.name = metadata["model_name"]
+
+ # Load relay IR.
+ relay_path = temp.relpath("model.json")
+ with open(relay_path, "r") as relay_file:
+ self.mod = tvm.ir.load_json(relay_file.read())
+
+ # Load parameter dictionary.
+ params_path = temp.relpath("model.params")
+ with open(params_path, "rb") as params_file:
+ self.params = relay.load_param_dict(params_file.read())
+
+ records_path = self.get_default_tuning_path()
+ if os.path.exists(records_path):
+ self.tuning_records = records_path
+
+ def get_default_tuning_path(self):
+ """Returns a default path for tuning records.
+
+ Returns
+ -------
+ tuning_records : str
+ A path in the models temporary directory that tuning
+ records can be stored.
+ """
+ return self._tmp_dir.relpath("tuning_records")
+
+ def export_package(
+ self,
+ executor_factory: GraphExecutorFactoryModule,
+ package_path: Optional[str] = None,
+ cross: Optional[Union[str, Callable]] = None,
+ ):
+ """Save this TVMCModel to file.
+
+ Parameters
+ ----------
+ executor_factory : GraphExecutorFactoryModule
+ The compiled graph factory that will be used to generate the
output package.
+ 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.
+ """
+ if package_path is None:
+ package_path = self._tmp_dir.relpath("package.tar")
+ export_model_library_format(executor_factory, package_path,
cross=cross)
+ self.package_path = package_path
+
+ def summary(self):
+ """ Print the IR corressponding to this model."""
+ print(self.mod)
+
+
+class TVMCPackage(object):
+ """Load a saved TVMCPackage from disk.
+
+ Parameters
+ ----------
+ package_path : str
+ The path to the saved TVMCPackage that will be loaded.
+ """
+
+ def __init__(self, package_path: str):
+ self._tmp_dir = utils.tempdir()
+ self.import_package(package_path)
+
+ def import_package(self, package_path: str):
+ """Load a TVMCPackage from a previously exported TVMCModel.
+
+ Parameters
+ ----------
+ package_path : str
+ The path to the saved TVMCPackage.
+ """
+ temp = self._tmp_dir
+ t = tarfile.open(package_path)
+ t.extractall(temp.relpath("."))
+
+ metadata_path = "metadata.json"
+ with open(temp.relpath(metadata_path), "r") as metadata_file:
+ metadata = json.load(metadata_file)
+ self.name = metadata["model_name"]
+ self.target = metadata["target"]
+ self.runtimes = metadata["runtimes"]
+
+ parameter_path = os.path.join("parameters", f"{self.name}.params")
+ with open(temp.relpath(parameter_path), "rb") as param_file:
+ self.params = bytearray(param_file.read())
+
+ graph_path = os.path.join("runtime-config", "graph", "graph.json")
+ with open(temp.relpath(graph_path), "r") as graph_file:
+ self.graph = graph_file.read()
+
+ ir_path = "relay.txt"
+ with open(temp.relpath(ir_path), "r") as relay_file:
+ self.mod = relay_file.read()
+
+ self.lib_path = temp.relpath(f"{self.name}.so")
+
+ def summary(self):
+ print(self.mod)
+
+
+class TVMCResult(object):
+ """Create a convenience wrapper around the output of tvmc.run
+
+ Parameters
+ ----------
+ outputs : dict
+ Outputs dictionary mapping the name of the output to its numpy value.
+ times : list of str
+ The execution times measured by the time evaluator to produce outputs.
+ """
+
+ def __init__(self, outputs: Dict[str, np.ndarray], times: List[str]):
+ self.outputs = outputs
+ self.times = times
+
+ def format_times(self):
+ """Format the mean, max, min and std of the execution times.
+
+ This has the effect of producing a small table that looks like:
+ .. code-block::
+ Execution time summary:
+ mean (ms) max (ms) min (ms) std (ms)
+ 0.14310 0.16161 0.12933 0.01004
+
+ Returns
+ -------
+ str
+ A formatted string containing the statistics.
+ """
+
+ # timestamps
+ mean_ts = np.mean(self.times) * 1000
+ std_ts = np.std(self.times) * 1000
+ max_ts = np.max(self.times) * 1000
+ min_ts = np.min(self.times) * 1000
+
+ header = "Execution time summary:\n{0:^10} {1:^10} {2:^10}
{3:^10}".format(
+ "mean (ms)", "max (ms)", "min (ms)", "std (ms)"
+ )
+ stats = "{0:^10.2f} {1:^10.2f} {2:^10.2f} {3:^10.2f}".format(
+ mean_ts, max_ts, min_ts, std_ts
+ )
+
+ return "%s\n%s\n" % (header, stats)
+
+ def get_top_results(self, max_results: int):
Review comment:
it seems like there is only going to be more introduced, though? I can't
see us deleting any. without providing a home for these, then the logical thing
to do in follow-on PRs will just be to add them here.
##########
File path: python/tvm/driver/tvmc/runner.py
##########
@@ -112,8 +112,10 @@ def drive_run(args):
except IOError as ex:
raise TVMCException("Error loading inputs file: %s" % ex)
- outputs, times = run_module(
- args.FILE,
+ tvmc_package = TVMCPackage(package_path=args.FILE)
Review comment:
@leandron It's a bit of a long explanation, I will see if I can
summarize. Currently a codegen backend is expected to produce runtime::Module.
When either the load mechanism (if an existing e.g. CSourceModule could
possibly be used) or the format (when producing e.g. not C) differs from one
checked-in, a new Module must be defined. For this reason, many codegen produce
aggregate modules containing multiple pieces of data. See CUDA for an example.
Worse yet, Modules are expected to implement their own load() and save(),
and it is possible (though not likely now) that you could produce a Module in
codegen which, after save() and load(), don't behave the same. It's possible we
are not entirely unit testing now what we deploy.
Finally, Modules aren't labelled--so while it kind of makes sense how we've
populated Model Library Format given target_host-targeted LLVM and C modules,
it's difficult to implement any sort of generic organizational scheme.
My opinion is that the solution to this problem is generating a set of POD
types (e.g. structs) from the compiler with labels and overhauling the module
load process to better align it with the implied load process we are asking
users to implement in microTVM. Specifically, it should be easy to document a
set of steps that the user needs to implement on their own in order to arrive
at an e.g. GraphExecutor from a given SoC after device reset.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
