huajsj commented on a change in pull request #9108:
URL: https://github.com/apache/tvm/pull/9108#discussion_r718011487
##########
File path: python/tvm/contrib/pipeline_executor.py
##########
@@ -533,11 +570,92 @@ def graph_executor_create(self, pipeline_mods,
mod_config):
The Modudle configuration.
"""
- mods = []
- for pipeline_mod in pipeline_mods:
- mod = graph_executor.GraphModule(
- pipeline_mod["default"](pipeline_mods[pipeline_mod]["dev"])
- )
- mods.append(mod.module)
+ # The module in pipeline_mods has a index information to identify the
modules order
+ # in pipe line, the pipe line executor follow such asend order to run
each module,
+ # but asend order of module is not guaranteed in pipeline_mod, here
need to pre-allocate
+ # module list then put module in correct place follow the index value.
+ mods = [None for _ in range(len(pipeline_mods))]
+ for lib_index in pipeline_mods:
+ pipeline_lib = pipeline_mods[lib_index]["lib"]
+ dev = pipeline_mods[lib_index]["dev"]
+ lib = graph_executor.GraphModule(pipeline_lib["default"](dev))
+ # Return a module list sorted by lib_index.
+ mods[lib_index - 1] = lib.module
return mods, json.dumps(mod_config)
+
+ def export_library(self, directory_path=None):
+ """Export pipeline runtime into disk.
+
+ Parameters
+ ----------
+ directory_path : str
+ The directory to which these files are exported.
+ """
+ if not self.pipeline_libs:
+ raise RuntimeError(f"The pipeline executor has not been
initialized.")
+
+ # If directory_path is not set, use the temporary path as the file
storage directory_path.
+ if not directory_path:
+ directory_path = tvm.contrib.utils.tempdir().temp_dir
+
+ # Create if the directory does not exist.
+ if not os.path.exists(directory_path):
+ os.makedirs(directory_path)
+ # Get the initial version of the configuration for export.
+ export_conf = self.mods_config.copy()
+ # Export library, JSON file and parameter file, and export the
corresponding
+ # relationship between the files and the pipeline modules.
+ for lib_index in self.pipeline_libs:
+ mconf = export_conf[lib_index - 1]
+ mconf["lib_name"] = "{}/lib{}.so".format(directory_path, lib_index)
+ mconf["json_name"] = "{}/json{}".format(directory_path, lib_index)
+ mconf["params_name"] = "{}/params{}".format(directory_path,
lib_index)
+ mconf["dev"] = "{},{}".format(
+ self.pipeline_libs[lib_index]["dev"].device_type,
+ self.pipeline_libs[lib_index]["dev"].device_id,
+ )
+
+ # Get graph, lib and parameters from GraphExecutorFactoryModule.
+ graph, lib, params = self.pipeline_libs[lib_index]["lib"]
+ # Export lib, graph and parameters to disk.
+ lib.export_library(mconf["lib_name"])
+ with open(mconf["json_name"], "w") as file_handle:
+ file_handle.write(graph)
+ with open(mconf["params_name"], "wb") as file_handle:
+ file_handle.write(relay.save_param_dict(params))
+
+ # Check whether the output is successful.
+ if not os.path.exists(mconf["json_name"]):
+ raise RuntimeError("File {} export
failure.".format(mconf["json_name"]))
+ if not os.path.exists(mconf["lib_name"]):
+ raise RuntimeError("File {} export
failure.".format(mconf["lib_name"]))
+ if not os.path.exists(mconf["params_name"]):
+ raise RuntimeError("File {} export
failure.".format(mconf["params_name"]))
Review comment:
removed.
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
