This is an automated email from the ASF dual-hosted git repository.
syfeng pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new e0518da2a5 [Unity][MSC][M2.3] Add tracker for track layer datas
(#16207)
e0518da2a5 is described below
commit e0518da2a597a61f5c5a6613b2a54b2e46a8039f
Author: Archermmt <[email protected]>
AuthorDate: Sun Dec 10 13:34:08 2023 +0800
[Unity][MSC][M2.3] Add tracker for track layer datas (#16207)
* add tracker
* remove distiller
* format fix
---
python/tvm/contrib/msc/core/ir/graph.py | 76 ++++++
python/tvm/contrib/msc/core/runtime/runner.py | 286 +++++++++++++++------
python/tvm/contrib/msc/core/tools/__init__.py | 1 +
python/tvm/contrib/msc/core/tools/prune/method.py | 6 +-
python/tvm/contrib/msc/core/tools/prune/pruner.py | 149 +++++++----
python/tvm/contrib/msc/core/tools/tool.py | 256 +++++++++---------
.../tools => core/tools/track}/__init__.py | 5 +-
python/tvm/contrib/msc/core/tools/track/method.py | 102 ++++++++
python/tvm/contrib/msc/core/tools/track/tracker.py | 185 +++++++++++++
python/tvm/contrib/msc/core/utils/info.py | 37 ++-
python/tvm/contrib/msc/core/utils/log.py | 21 +-
python/tvm/contrib/msc/core/utils/message.py | 23 ++
.../msc/framework/tensorflow/runtime/runner.py | 2 +-
.../msc/framework/tensorflow/tools/__init__.py | 1 +
.../tensorflow/tools/{ => track}/__init__.py | 4 +-
.../framework/tensorflow/tools/track/tracker.py | 55 ++++
.../msc/framework/tensorrt/tools/__init__.py | 1 +
.../tools => tensorrt/tools/track}/__init__.py | 4 +-
.../msc/framework/tensorrt/tools/track/tracker.py | 159 ++++++++++++
.../msc/framework/torch/frontend/translate.py | 2 +-
.../contrib/msc/framework/torch/tools/__init__.py | 1 +
.../tools => torch/tools/track}/__init__.py | 4 +-
.../msc/framework/torch/tools/track/tracker.py | 55 ++++
.../contrib/msc/framework/tvm/tools/__init__.py | 1 +
.../tools => tvm/tools/track}/__init__.py | 4 +-
.../msc/framework/tvm/tools/track/tracker.py | 155 +++++++++++
python/tvm/contrib/msc/pipeline/manager.py | 112 +++++---
src/contrib/msc/core/ir/graph.cc | 17 ++
src/contrib/msc/core/ir/graph.h | 2 +-
src/contrib/msc/framework/tensorrt/codegen.cc | 37 ++-
src/contrib/msc/framework/tensorrt/codegen.h | 4 +
tests/python/contrib/test_msc/test_tools.py | 102 +++++---
32 files changed, 1491 insertions(+), 378 deletions(-)
diff --git a/python/tvm/contrib/msc/core/ir/graph.py
b/python/tvm/contrib/msc/core/ir/graph.py
index 8fabed30ac..51727cd089 100644
--- a/python/tvm/contrib/msc/core/ir/graph.py
+++ b/python/tvm/contrib/msc/core/ir/graph.py
@@ -104,6 +104,17 @@ class MSCTensor(Object):
return False
return True
+ def to_json(self) -> str:
+ """Dump the tensor to json.
+
+ Returns
+ -------
+ tensor_json: string
+ The tensor in json format.
+ """
+
+ return _ffi_api.MSCTensorToJson(self)
+
def inspect(self) -> dict:
"""Extract important info of the tensor.
@@ -117,6 +128,45 @@ class MSCTensor(Object):
tensor_des["layout"] = self.layout.name if self.layout else ""
return tensor_des
+ @classmethod
+ def from_json(cls, json_str: str, **options) -> object:
+ """Load the tensor from json.
+
+ Parameters
+ ----------
+ json_str: string
+ The file_path or json string.
+ options: dict
+ The items to be changed.
+
+ Returns
+ -------
+ tensor: MSCTensor
+ The tensor.
+ """
+
+ dict_obj = msc_utils.load_dict(json_str)
+ dict_obj.update(options)
+ return _ffi_api.MSCTensorFromJson(msc_utils.dump_dict(dict_obj))
+
+ def clone(self, **options) -> object:
+ """Clone the tensor.
+
+ Parameters
+ ----------
+ json_str: string
+ The file_path or json string.
+ options: dict
+ The items to be changed.
+
+ Returns
+ -------
+ new_tensor: MSCTensor
+ The cloned tensor.
+ """
+
+ return MSCTensor.from_json(self.to_json(), **options)
+
@property
def dtype_name(self) -> str:
return _ffi_api.MSCTensorDTypeName(self)
@@ -391,6 +441,19 @@ class WeightJoint(BaseJoint):
friends,
)
+ def set_attr(self, key: str, value: str):
+ """Set attribute to node
+
+ Parameters
+ -------
+ key: str
+ The key of the attribute.
+ value: str
+ The value.
+ """
+
+ _ffi_api.WeightJointSetAttr(self, key, value)
+
def get_attrs(self) -> Dict[str, str]:
"""Get all the attributes from node
@@ -536,6 +599,19 @@ class MSCGraph(BaseGraph):
return _ffi_api.MSCGraphFindTensor(self, name)
+ def set_tensor_alias(self, tensor: MSCTensor, alias: str):
+ """Set alis for the tensor
+
+ Parameters
+ -------
+ tensor: MSCTensor
+ The tensor.
+ alias: str
+ The alias.
+ """
+
+ _ffi_api.MSCGraphSetTensorAlias(self, tensor, alias)
+
def find_producer(self, ref: Union[str, MSCTensor]) -> MSCJoint:
"""Find producer by tensor_name or tensor.
diff --git a/python/tvm/contrib/msc/core/runtime/runner.py
b/python/tvm/contrib/msc/core/runtime/runner.py
index 12e2dec02a..dcf24225fe 100644
--- a/python/tvm/contrib/msc/core/runtime/runner.py
+++ b/python/tvm/contrib/msc/core/runtime/runner.py
@@ -20,13 +20,13 @@
import os
import json
import logging
-from typing import Dict, Optional, Any, List, Tuple, Union
+from typing import Dict, Optional, Any, List, Tuple, Union, Iterable
import numpy as np
import tvm
from tvm.contrib.msc.core.ir import MSCGraph
from tvm.contrib.msc.core.frontend import from_relax
-from tvm.contrib.msc.core.tools import BaseTool, ToolType, create_tool
+from tvm.contrib.msc.core.tools import BaseTool, ToolType, create_tool,
remove_tools
from tvm.contrib.msc.core.utils.namespace import MSCFramework
from tvm.contrib.msc.core.utils.message import MSCStage
from tvm.contrib.msc.core import utils as msc_utils
@@ -76,9 +76,9 @@ class BaseRunner(object):
logger: logging.Logger = None,
):
self._mod = mod
- self._tools_config = tools_config or {}
- self._translate_config = translate_config or {}
- self._generate_config = generate_config or {}
+ self._tools_config = msc_utils.copy_dict(tools_config)
+ self._translate_config = msc_utils.copy_dict(translate_config)
+ self._generate_config = msc_utils.copy_dict(generate_config)
self._stage = stage
self._name = name
self._device = device if self._device_enabled(device) else "cpu"
@@ -102,16 +102,19 @@ class BaseRunner(object):
if "build_folder" not in self._generate_config:
self._generate_config["build_folder"] = msc_utils.get_build_dir()
- if self._tools_config:
- if "codegen" not in self._generate_config:
- self._generate_config["codegen"] = {}
- self._generate_config["codegen"].update({"use_tools": True,
"tools_tag": self._name})
self._graphs, self._weights = [], []
self._model, self._model_info = None, {}
self._runnable = None
+ # Setup tools
self._tools = {}
+ if self._tools_config:
+ self._update_codegen({"use_tools": True, "tools_tag": self._name})
+ for t_type, config in self._tools_config.items():
+ self._tools[t_type] = create_tool(
+ self.framework, t_type, self._name, stage=self._stage,
**config
+ )
return {
- "tools_config": self._tools_config,
+ "tools": {k: v.tool_style() for k, v in self._tools.items()},
"translate_config": self._translate_config,
"generate_config": self._generate_config,
"name": self._name,
@@ -121,21 +124,28 @@ class BaseRunner(object):
}
def change_stage(self, stage: str):
- """Change the stage of tools and strategy"""
+ """Change the stage of runner and tools"""
self._stage = stage
for tool in self._tools.values():
tool.change_stage(stage)
- def build(self, cache_dir: msc_utils.MSCDirectory = None, build_graph:
bool = False) -> Any:
+ def change_logger(self, logger: logging.Logger):
+ """Change the logger of runner and tools"""
+
+ self._logger = logger
+ for tool in self._tools.values():
+ tool.change_logger(logger)
+
+ def build(self, cache_dir: msc_utils.MSCDirectory = None, force_build:
bool = False) -> Any:
"""Build the runnable object
Parameters
-------
cache_dir: MSCDirectory
cache path for save/load info
- build_graph: bool
- Whether to build the MSCGraphs.
+ force_build: bool
+ Whether to force build the runner.
Returns
-------
@@ -143,89 +153,137 @@ class BaseRunner(object):
The runnable object.
"""
+ if force_build:
+ self._graphs, self._weights = [], []
+ self._model, self._model_info = None, {}
+ self._runnable = None
if cache_dir and os.path.isfile(cache_dir.relpath("cache_info.json")):
cache_info =
msc_utils.load_dict(cache_dir.relpath("cache_info.json"))
else:
cache_info = {}
- # Create tools
- if self._tools_config:
- for t_type, config in self._tools_config.items():
- self._tools[t_type] = create_tool(
- self.framework, t_type, self._name, stage=self._stage,
**config
- )
-
# Load graphs from cache
- if cache_info.get("graphs"):
+ if not self._graphs and cache_info.get("graphs"):
self._graphs, self._weights = self._load_graphs(cache_dir,
cache_info["graphs"])
- self._logger.debug(
- "Load {} graphs from cache @ {}".format(len(self._graphs),
cache_dir)
- )
+ self._logger.info("Load %d graphs from %s", len(self._graphs),
cache_dir)
- # Get or rebuild graphs
- if build_graph or not self._graphs:
+ # Translate graphs from module
+ if not self._graphs:
self._graphs, self._weights = self._translate()
- self._logger.debug("Translate {} graphs from
module".format(len(self._graphs)))
+ self._logger.info("Translate %d graphs from module",
len(self._graphs))
- # reset graph for tools
- for tool in self._tools.values():
- self._graphs, self._weights = tool.reset(
- self._graphs, self._weights, cache_dir=cache_dir
- )
-
- if cache_info.get("model") and not build_graph:
- # Load model from cache
+ # Load model from cache
+ if not self._model and cache_info.get("model"):
+ self._graphs, self._weights = self.reset_tools(cache_dir=cache_dir)
self._model = self._load_model(cache_dir, cache_info["model"])
- else:
- # Generate model
+ self._logger.info("Load model(%s) from %s", self.framework,
cache_dir)
+
+ # Generate model
+ if not self._model:
+ # Generate normal model
+ self._graphs, self._weights = self.reset_tools(cache_dir=cache_dir)
self._model = self._generate_model()
- if "loader" in self._generate_config:
- loader, generate_config = self._generate_config["loader"]
- self._model = loader(self._model, **generate_config)
- self._logger.info(
- "Model({}) processed by customize loader {}({})".format(
- self.framework, loader, generate_config
- )
- )
+
+ # Log generate info
+ generate_msg = "Generate model({})".format(self.framework)
+ if self._tools:
+ self._logger.info("%s with tools: %s", generate_msg,
",".join(self._tools.keys()))
+ else:
+ self._logger.info("%s without tools", generate_msg)
+ if "generator" in self._generate_config:
+ generator, generate_config = self._generate_config["generator"]
+ self._model = generator(self._model, **generate_config)
+ self._logger.info("%s by %s(%s)", generate_msg, generator,
generate_config)
+
+ # Inspect model
self._model_info = self._inspect_model()
- if self._debug_level >= 3:
+ if self._debug_level >= 2:
self._logger.debug(msc_utils.msg_block("RUNNER.MODEL_INFO",
self._model_info))
- if cache_info.get("runnable") and not build_graph:
- # Load runnable from cache
+ runnable_msg = "runnable({}, {}) @ {}".format(
+ self.framework, "train" if self._is_training else "eval",
self._device
+ )
+
+ # Load runnable from cache
+ if not self._runnable and cache_info.get("runnable"):
self._runnable = self._load_runnable(cache_dir,
cache_info["runnable"])
- else:
- # Build runnable on device
+ self._logger.info("Load %s from %s", runnable_msg, cache_dir)
+
+ # Build runnable
+ if not self._runnable:
self._runnable = self._to_runnable(self._model, self._device,
self._is_training)
- self._logger.info(
- "Runnable({}, {}) loaded on device {}".format(
- self.framework, "train" if self._is_training else "eval",
self._device
- )
- )
+ self._logger.info("Build %s", runnable_msg)
return self._runnable
- def save_cache(self, cache_dir: msc_utils.MSCDirectory):
+ def save_cache(
+ self,
+ cache_dir: msc_utils.MSCDirectory,
+ save_model: bool = True,
+ save_runnable: bool = True,
+ save_tools: bool = True,
+ ):
"""Save runner to cache
Parameters
-------
cache_dir: MSCDirectory
cache path for save/load info
+ save_model: bool
+ Whether to save model.
+ save_runnable: bool
+ Whether to save runnable.
+ save_tools: bool
+ Whether to save tools.
"""
- cache_info = {
- "graphs": self._save_graphs(cache_dir),
- "model": self._save_model(cache_dir),
- "runnable": self._save_runnable(cache_dir),
- }
- for tool in self._tools.values():
- cache_info.update(tool.save_cache(cache_dir))
+ cache_info = {"graphs": self._save_graphs(cache_dir)}
+ if save_model:
+ cache_info["model"] = self._save_model(cache_dir)
+ if save_runnable:
+ cache_info["runnable"] = self._save_runnable(cache_dir)
+ if save_tools:
+ for t_type, tool in self._tools.items():
+ cache_info[t_type] = tool.save_cache(cache_dir)
with open(cache_dir.relpath("cache_info.json"), "w") as f:
f.write(json.dumps(cache_info, indent=2))
- if self._debug_level >= 3:
- self._logger.debug(
- msc_utils.msg_block("RUNNER.CACHE_INFO", {"folder": cache_dir,
"info": cache_info})
- )
+ self._logger.debug(
+ msc_utils.msg_block("RUNNER.SAVE_CACHE", {"folder": cache_dir,
"info": cache_info})
+ )
+
+ def reset_tools(
+ self,
+ graphs: List[MSCGraph] = None,
+ weights: List[Dict[str, tvm.nd.array]] = None,
+ tools: List[BaseTool] = None,
+ cache_dir: msc_utils.MSCDirectory = None,
+ ):
+ """Reset the tools
+
+ Parameters
+ -------
+ graphs: list<MSCgraph>
+ The msc graphs.
+ weights: list<dict<str, tvm.nd.array>>
+ The weights.
+ tools: list<BaseTool>
+ The tools.
+ cache_dir: MSCDirectory
+ cache path for save/load info.
+
+ Returns
+ -------
+ graphs: list<MSCgraph>
+ The msc graphs.
+ weights: list<dict<str, tvm.nd.array>>
+ The weights.
+ """
+
+ graphs = graphs or self._graphs
+ weights = weights or self._weights
+ tools = tools or self._tools.values()
+ for tool in tools:
+ graphs, weights = tool.reset(graphs, weights, cache_dir)
+ return graphs, weights
def run(
self, inputs: Union[List[np.ndarray], Dict[str, np.ndarray]],
ret_type="dict"
@@ -260,6 +318,8 @@ class BaseRunner(object):
), "Expected all inputs as np.ndarray"
inputs = {i["name"]: inputs[i["name"]] for i in model_inputs}
outputs = self._call_runnable(self._runnable, inputs, self._device)
+ if ret_type == "native":
+ return outputs
if ret_type == "dict":
if isinstance(outputs, (list, tuple)):
assert len(outputs) == len(
@@ -283,6 +343,22 @@ class BaseRunner(object):
outputs = [msc_utils.cast_array(data) for data in outputs]
return outputs
+ def get_tool_config(self, tool_type: str) -> dict:
+ """Get tool by type
+
+ Parameters
+ -------
+ tool_type: str
+ The type of the tool prune| quantize| distill...
+
+ Returns
+ -------
+ config: dict
+ The tool config.
+ """
+
+ return self._tools_config.get(tool_type)
+
def get_tool(self, tool_type: str) -> BaseTool:
"""Get tool by type
@@ -299,7 +375,21 @@ class BaseRunner(object):
return self._tools.get(tool_type)
- def apply_tool(self, tool_type: str, data_loader: Any = None) -> dict:
+ def get_tools(self) -> Iterable[BaseTool]:
+ """Get all saved tools by tag
+
+ Returns
+ -------
+ tools: iterable<BaseTool>
+ The saved tools.
+ """
+
+ for t_type in ToolType.all_types():
+ tool = self.get_tool(t_type)
+ if tool:
+ yield tool
+
+ def apply_tool(self, tool_type: str, data_loader: Any = None) -> str:
"""Execute tool and get plan
Parameters
@@ -308,6 +398,11 @@ class BaseRunner(object):
The tool type, should be in ToolType
data_loader:
The data loader
+
+ Returns
+ -------
+ plan_file: str
+ The saved plan file.
"""
assert tool_type in self._tools, "Can not find tool " + str(tool_type)
@@ -316,7 +411,7 @@ class BaseRunner(object):
if not pruner.finalize():
assert data_loader, "data_loader should be given to plan prune"
for inputs in data_loader():
- self.run(inputs)
+ self.run(inputs, ret_type="native")
break
plan = pruner.finalize()
else:
@@ -325,8 +420,28 @@ class BaseRunner(object):
plan_file = self._tools_config[tool_type]["plan_file"]
with open(plan_file, "w") as f:
f.write(json.dumps(plan, indent=2))
- self._logger.info("Save %s plan -> %s", tool_type, plan_file)
- return plan
+ self._logger.info("Save %d plan(%s) -> %s", len(plan), tool_type,
plan_file)
+ return plan_file
+
+ def _update_codegen(self, config: Dict[str, Any]):
+ """Update the codegen in generate_config
+
+ Parameters
+ -------
+ config: dict
+ The extra config for codegen.
+ """
+
+ if "codegen" not in self._generate_config:
+ self._generate_config["codegen"] = {}
+ codegen = self._generate_config["codegen"]
+ if isinstance(codegen, dict):
+ codegen.update(config)
+ elif isinstance(codegen, (list, tuple)):
+ for c in codegen:
+ c.update(config)
+ else:
+ raise TypeError("Unexpecet codegen config " + str(codegen))
def visualize(self, visual_dir: msc_utils.MSCDirectory):
"""Visualize MSCGraphs
@@ -371,8 +486,9 @@ class BaseRunner(object):
self._model = None
if self._runnable:
self._runnable = None
- for tool in self._tools.values():
+ for tool in self.get_tools():
tool.destory()
+ remove_tools(self._name)
def _translate(self) -> Tuple[List[MSCGraph], Dict[str, tvm.nd.array]]:
"""Translate IRModule to MSCgraphs
@@ -434,7 +550,7 @@ class BaseRunner(object):
-------
graphs: list<MSCgraph>
The msc graphs.
- weights: list<dic<str, tvm.nd.array>>
+ weights: list<dict<str, tvm.nd.array>>
The weights
Returns
@@ -579,10 +695,18 @@ class BaseRunner(object):
return True
+ @classmethod
+ def support_tool(cls, tool_type: str) -> bool:
+ return True
+
@property
def stage(self):
return self._stage
+ @property
+ def debug_level(self):
+ return self._debug_level
+
@property
def model(self):
return self._model
@@ -690,7 +814,7 @@ class ModelRunner(BaseRunner):
-------
graphs: list<MSCgraph>
The msc graphs.
- weights: list<dic<str, tvm.nd.array>>
+ weights: list<dict<str, tvm.nd.array>>
The weights
Returns
@@ -699,9 +823,11 @@ class ModelRunner(BaseRunner):
The runnable model
"""
+ graph = graphs[0] if graphs else self._graphs[0]
+ weight = weights[0] if weights else self._weights[0]
return self.codegen_func(
- graphs or self._graphs[0],
- weights or self._weights[0],
+ graph,
+ weight,
codegen_config=self._generate_config.get("codegen"),
print_config=self._generate_config.get("print"),
build_folder=self._generate_config["build_folder"],
@@ -853,7 +979,7 @@ class BYOCRunner(BaseRunner):
-------
graphs: list<MSCgraph>
The msc graphs.
- weights: list<dic<str, tvm.nd.array>>
+ weights: list<dict<str, tvm.nd.array>>
The weights
Returns
@@ -954,10 +1080,10 @@ class BYOCRunner(BaseRunner):
The inspected model info
"""
- if self._debug_level >= 3:
+ if self._debug_level >= 2:
for idx, graph in enumerate(self._graphs):
self._logger.debug(
- msc_utils.msg_block("RUNNER.GRAPH[{}].INFO".format(idx),
graph.inspect())
+ msc_utils.msg_block("GRAPH[{}].INFO".format(idx),
graph.inspect())
)
return self._byoc_graph.inspect()
diff --git a/python/tvm/contrib/msc/core/tools/__init__.py
b/python/tvm/contrib/msc/core/tools/__init__.py
index 3d60ce22c6..0524e4c823 100644
--- a/python/tvm/contrib/msc/core/tools/__init__.py
+++ b/python/tvm/contrib/msc/core/tools/__init__.py
@@ -19,3 +19,4 @@
from .tool import *
from .execute import *
from .prune import *
+from .track import *
diff --git a/python/tvm/contrib/msc/core/tools/prune/method.py
b/python/tvm/contrib/msc/core/tools/prune/method.py
index 6e294dc15c..fd3abe8df4 100644
--- a/python/tvm/contrib/msc/core/tools/prune/method.py
+++ b/python/tvm/contrib/msc/core/tools/prune/method.py
@@ -99,9 +99,9 @@ class PruneMethod(object):
return config
if len(in_indices) > 0:
data = cls.prune_axis(data, in_axis, in_indices)
- weight = pruner.find_tensor(name)
- left_num = int(((density * weight.dim_at(out_axis) + stride) //
stride) * stride)
- axis_sum = [np.abs(d).sum() for d in np.split(data,
data.shape[out_axis], out_axis)]
+ out_dim = data.shape[out_axis]
+ left_num = int(((density * out_dim + stride) // stride) * stride)
+ axis_sum = [np.abs(d).sum() for d in np.split(data, out_dim, out_axis)]
rank = np.argsort(np.array(axis_sum))
config["out_indices"] = rank[-left_num:].tolist()
return config
diff --git a/python/tvm/contrib/msc/core/tools/prune/pruner.py
b/python/tvm/contrib/msc/core/tools/prune/pruner.py
index 5fa625c2ad..577fe2bf4f 100644
--- a/python/tvm/contrib/msc/core/tools/prune/pruner.py
+++ b/python/tvm/contrib/msc/core/tools/prune/pruner.py
@@ -17,10 +17,11 @@
"""tvm.contrib.msc.core.tools.prune.pruner"""
from typing import List, Dict, Tuple, Any
+import numpy as np
import tvm
from tvm.contrib.msc.core.ir import MSCGraph, WeightJoint, MSCTensor
-from tvm.contrib.msc.core.tools.tool import ToolType, WeightTool, Strategy
+from tvm.contrib.msc.core.tools.tool import ToolType, WeightTool, ToolStrategy
from tvm.contrib.msc.core import _ffi_api
from tvm.contrib.msc.core import utils as msc_utils
from .method import PruneMethod
@@ -64,7 +65,7 @@ class BasePruner(WeightTool):
}
return main_wtypes, relation_wtypes
- def _parse_strategys(self, strategy_list: dict) -> Dict[str, Strategy]:
+ def _parse_strategys(self, strategy_list: dict) -> Dict[str, ToolStrategy]:
"""Parse the strategy to get valid strategy
Parameters
@@ -74,7 +75,7 @@ class BasePruner(WeightTool):
Returns
-------
- strategys: dict<str, Strategy>
+ strategys: dict<str, ToolStrategy>
The parsed strategy.
"""
@@ -85,10 +86,10 @@ class BasePruner(WeightTool):
return super()._parse_strategys([_update_stages(s) for s in
strategy_list])
- def load_graphs(
+ def _reset(
self, graphs: List[MSCGraph], weights: List[Dict[str, tvm.nd.array]]
) -> Tuple[List[MSCGraph], List[Dict[str, tvm.nd.array]]]:
- """Load the graphs and weights
+ """Reset the tool
Parameters
----------
@@ -105,10 +106,13 @@ class BasePruner(WeightTool):
The weights
"""
- graphs, weights = super().load_graphs(graphs, weights)
- if not self._plan:
- return graphs, weights
- return self.prune_graphs(graphs, weights)
+ self._meta_weights = {}
+ for sub_weights in weights:
+ self._meta_weights.update(sub_weights)
+ graphs, weights = super()._reset(graphs, weights)
+ if self._plan and self._enabled:
+ return self.prune_graphs(graphs, weights)
+ return graphs, weights
def _execute_before_build(self, *args, **kwargs):
"""Execute before model build
@@ -164,12 +168,10 @@ class BasePruner(WeightTool):
strategy = self._get_tensor_strategy(name, consumer)
if not strategy:
return False
- if strategy.get_config("density", 1.0) == 1.0:
- return False
return True
def _process_tensor(
- self, tensor: Any, name: str, consumer: str, scope: str, strategys:
List[Strategy]
+ self, tensor: Any, name: str, consumer: str, scope: str, strategys:
List[ToolStrategy]
) -> Any:
"""Process tensor
@@ -183,7 +185,7 @@ class BasePruner(WeightTool):
The name of the consumer.
scope: str
The scope mark teacher| student| null.
- strategys: list<Strategy>
+ strategys: list<ToolStrategy>
The strategys for the tensor.
Returns
@@ -215,7 +217,7 @@ class BasePruner(WeightTool):
}
return tensor
- def _prune_tensor(self, name: str, consumer: str, strategys:
List[Strategy]) -> Any:
+ def _prune_tensor(self, name: str, consumer: str, strategys:
List[ToolStrategy]) -> Any:
"""Prune tensor
Parameters
@@ -226,7 +228,7 @@ class BasePruner(WeightTool):
The name of the consumer.
scope: str
The scope mark teacher| student| null.
- strategys: list<Strategy>
+ strategys: list<ToolStrategy>
The strategys for the tensor.
"""
@@ -250,6 +252,8 @@ class BasePruner(WeightTool):
def _prunable(w_node: WeightJoint) -> bool:
"""Check if weight node is prunable"""
+ if strategy.get_config().get("density", 1) == 1:
+ return False
if w_node.get_attr("weight_strategy") != "main":
return False
if not w_node.children:
@@ -283,7 +287,7 @@ class BasePruner(WeightTool):
elif _prunable(w_node):
self._plan[w_node.name] = strategy(
self,
- self.get_data(w_node.name),
+ self.get_meta_data(w_node.name),
w_node.name,
consumer,
in_axis=in_axis,
@@ -333,8 +337,10 @@ class BasePruner(WeightTool):
pruned_tensors, pruned_weights = {}, {}
for node in graph.get_nodes():
for weight in node.get_weights().values():
- w_name = weight.name
- if w_name in self._plan and not
self._plan[w_name].get("pruned", False):
+ w_node, w_name = self.find_w_node(weight.name), weight.name
+ if w_name not in self._plan or w_node.get_attr("status",
"") == "pruned":
+ pruned_weights[w_name] = sub_weights[w_name]
+ else:
data = msc_utils.cast_array(sub_weights[w_name])
in_axis, out_axis =
self._get_io_axes(self.find_w_node(w_name))
w_config = self._plan[w_name]
@@ -344,23 +350,18 @@ class BasePruner(WeightTool):
data = PruneMethod.prune_axis(data, out_axis,
w_config["out_indices"])
pruned_tensors[w_name] = _prune_by_shape(weight,
data.shape)
pruned_weights[w_name] = tvm.nd.array(data)
- self._plan[w_name]["pruned"] = True
+ w_node.set_attr("status", "pruned")
pruned_weights_cnt += 1
- else:
- pruned_weights[w_name] = sub_weights[w_name]
- if node.optype == "constant" and node.weight_at("const").name
in pruned_tensors:
+ if node.optype == "constant":
+ if node.weight_at("const").name not in pruned_tensors:
+ continue
ref_tensor = pruned_tensors[node.weight_at("const").name]
- pruned_tensors[node.output_at(0).name] = MSCTensor(
- node.output_at(0).name,
- ref_tensor.dtype,
- ref_tensor.layout.name,
- ref_tensor.get_shape(),
- ref_tensor.alias,
+ pruned_tensors[node.output_at(0).name] = ref_tensor.clone(
+ name=node.output_at(0).name
)
- elif (
- node.optype in ("nn.conv2d", "msc.conv2d_bias",
"msc.linear", "msc.linear_bias")
- and node.weight_at("weight").name in pruned_tensors
- ):
+ elif node.optype in self._main_wtypes:
+ if node.weight_at("weight").name not in pruned_tensors:
+ continue
out = node.output_at(0)
if node.optype in ("msc.linear", "msc.linear_bias"):
channel_axis = out.ndim - 1
@@ -371,48 +372,82 @@ class BasePruner(WeightTool):
pruned_tensors[node.weight_at("weight").name].dim_at("O"),
channel_axis,
)
- else:
+ elif node.optype in self._relation_wtypes:
for out in node.get_outputs():
- if out.name in self._plan and not
self._plan[out.name].get("pruned", False):
- pruned_tensors[out.name] = _prune_by_channel(
- out, len(self._plan[out.name]["out_indices"])
- )
- self._plan[out.name]["pruned"] = True
- elif (
- node.get_inputs()
- and node.input_at(0).name in pruned_tensors
- and node.input_at(0).layout_of("C") >= 0
- and out.layout_of("C") >= 0
- ):
- pruned_tensors[out.name] = _prune_by_channel(
- out,
pruned_tensors[node.input_at(0).name].dim_at("C")
- )
+ w_node = self.find_w_node(out.name)
+ if out.name not in self._plan or
w_node.get_attr("status", "") == "pruned":
+ continue
+ pruned_tensors[out.name] = _prune_by_channel(
+ out, len(self._plan[out.name]["out_indices"])
+ )
+ w_node.set_attr("status", "pruned")
+ elif node.get_inputs():
+ ref_input = node.input_at(0)
+ if ref_input.name not in pruned_tensors or
ref_input.layout_of("C") < 0:
+ continue
+ for out in node.get_outputs():
+ if out.layout_of("C") < 0:
+ continue
+ pruned_tensors[out.name] = _prune_by_channel(
+ out, pruned_tensors[ref_input.name].dim_at("C")
+ )
+
+ def _is_pruned(tensor: MSCTensor, graph: MSCGraph) -> bool:
+ return tensor.get_shape() !=
graph.find_tensor(tensor.name).get_shape()
+
+ pruned_tensors = {k: v for k, v in pruned_tensors.items() if
_is_pruned(v, graph)}
if self.on_debug(3, in_forward=False):
self._logger.debug(msc_utils.msg_block("Pruned Tensors",
pruned_tensors))
- pruned_graph = _ffi_api.PruneWeights(graph, pruned_tensors)
- new_graphs.append(pruned_graph)
+
+ if pruned_tensors:
+ pruned_graph = _ffi_api.PruneWeights(graph, pruned_tensors)
+ new_graphs.append(pruned_graph)
+ else:
+ new_graphs.append(graph)
new_weights.append(pruned_weights)
- # log compress rate
def _flatten_size(weights):
weight_size = 0
for sub_weights in weights:
for w_data in sub_weights.values():
weight_size += w_data.asnumpy().size
- return weight_size
+ return weight_size / 2**20
raw_size = _flatten_size(weights)
- new_size = _flatten_size(new_weights)
- self._logger.info(
- "Prune {} weights, compress to {:g}% ({:g} M->{:g} M)".format(
+ # log compress rate
+ if pruned_weights_cnt > 0:
+ new_size = _flatten_size(new_weights)
+ self._logger.info(
+ "Prune %d weights, compress to %.2f%% (%.4f M->%.4f M)",
pruned_weights_cnt,
new_size * 100 / raw_size,
- raw_size / 2**20,
- new_size / 2**20,
+ raw_size,
+ new_size,
)
- )
+ else:
+ self._logger.info("No weights pruned, size %.4f M", raw_size)
return new_graphs, new_weights
+ def get_meta_data(self, name: str) -> np.ndarray:
+ """Get meta weight as np.ndarray
+
+ Parameters
+ ----------
+ name: str
+ The name of data.
+
+ Returns
+ -------
+ data: np.ndarray
+ The data in np.ndarray format.
+ """
+
+ if name in self._meta_weights:
+ return msc_utils.cast_array(self._meta_weights[name])
+ raise Exception(
+ "Can not find data {} from {} weights".format(name,
len(self._meta_weights))
+ )
+
def finalize(self) -> dict:
"""Get the plan"""
diff --git a/python/tvm/contrib/msc/core/tools/tool.py
b/python/tvm/contrib/msc/core/tools/tool.py
index c37ec3db97..480705f31b 100644
--- a/python/tvm/contrib/msc/core/tools/tool.py
+++ b/python/tvm/contrib/msc/core/tools/tool.py
@@ -54,7 +54,7 @@ class ToolScope(object):
STUDENT = "student"
-class Executor(object):
+class ToolExecutor(object):
"""Executor for process the tensor
Parameters
@@ -94,11 +94,6 @@ class Executor(object):
kwargs.update({k: v for k, v in self._config.items() if k not in
kwargs})
return self._method(*args, **kwargs)
- def get_config(self, key: str, default: Any) -> Any:
- """Get the value in config"""
-
- return self._config.get(key, default)
-
def copy(self, name: str = None, method: callable = None, config: dict =
None):
"""Copy a executor
@@ -113,20 +108,24 @@ class Executor(object):
Returns
-------
- new_strategy: Strategy
- The copied strategy
+ new_executor: ToolExecutor
+ The copied executor
"""
new_config = config or {}
new_config.update({k: v for k, v in self._config.items() if k not in
new_config})
- return Executor(name or self._name, method or self._method, new_config)
+ return ToolExecutor(name or self._name, method or self._method,
new_config)
@property
def name(self):
return self._name
+ @property
+ def config(self):
+ return self._config
+
-class Strategy(object):
+class ToolStrategy(object):
"""Strategy for process tensor
Parameters
@@ -137,13 +136,16 @@ class Strategy(object):
The tensor type.
stage: str
The init stage
+ meta: dict:
+ The meta strategy config.
"""
- def __init__(self, name: str, tensor_type: str, stage: str = "default"):
+ def __init__(self, name: str, tensor_type: str, stage: str = "default",
meta: dict = None):
self._name = name
self._tensor_type = tensor_type
self._stage = stage
self._executors = {}
+ self._meta = meta
def __str__(self):
return "{}({} @ {}) ".format(self._name, self._tensor_type,
self._stage) + "; ".join(
@@ -187,14 +189,14 @@ class Strategy(object):
self._stage = stage
- def add_executor(self, stage: str, executor: Executor):
+ def add_executor(self, stage: str, executor: ToolExecutor):
"""Add a executor to strategy
Parameters
----------
stage: str
The mark of the executor.
- executor: Executor
+ executor: ToolExecutor
The executor to process tensor.
"""
@@ -215,10 +217,10 @@ class Strategy(object):
return self._executors[self._stage]
return self._executors["default"]
- def get_config(self, key: str, default: Any) -> Any:
- """Get the value in config"""
+ def get_config(self) -> dict:
+ """Get the config of current executor"""
- return self.get_executor().get_config(key, default)
+ return self.get_executor().config
def support_stage(self, stage: str) -> bool:
"""Check if the strategy support a stage
@@ -258,12 +260,12 @@ class Strategy(object):
Returns
-------
- new_strategy: Strategy
+ new_strategy: ToolStrategy
The copied strategy
"""
configs = configs or {}
- strategy = Strategy(
+ strategy = ToolStrategy(
name or self._name, tensor_type or self._tensor_type, stage or
self._stage
)
for st_name, executor in self._executors.items():
@@ -271,6 +273,10 @@ class Strategy(object):
strategy.add_executor(st_name, new_executor)
return strategy
+ @property
+ def meta(self):
+ return self._meta
+
class BaseTool(object):
"""Basic tool of MSC
@@ -299,7 +305,7 @@ class BaseTool(object):
self,
stage: str,
plan_file: str,
- strategys: dict,
+ strategys: List[dict],
cache_processed: bool = True,
options: dict = None,
debug_level: int = 0,
@@ -320,9 +326,8 @@ class BaseTool(object):
title = "{}.SETUP({} @ {})".format(self.tool_type().upper(),
self._stage, self.framework())
self._logger.info(msc_utils.msg_block(title, self.setup(), width=0))
if self._debug_level >= 3 and self._plan:
- self._logger.debug(
-
msc_utils.msg_block("{}.PLAN".format(self.tool_type().upper()), self._plan)
- )
+ title = "{}.PLAN".format(self.tool_type().upper())
+ self._logger.debug(msc_utils.msg_block(title, self._plan))
def setup(self) -> dict:
"""Setup the tool
@@ -348,58 +353,62 @@ class BaseTool(object):
"debug_level": self._debug_level,
}
- def _parse_strategys(self, strategy_list: dict) -> Dict[str, Strategy]:
+ def _parse_strategys(self, strategy_list: List[dict]) -> Dict[str,
ToolStrategy]:
"""Parse the strategy to get valid strategy
Parameters
-------
- strategy_list: dict
- The given strategy
+ strategy_list: list<dict>
+ The given strategys
Returns
-------
- strategys: dict<str, Strategy>
+ strategys: dict<str, ToolStrategy>
The parsed strategy.
"""
strategys = {}
assert isinstance(strategy_list, list) and all(
isinstance(s, dict) for s in strategy_list
- ), "Strategy should be given as list of dict"
- for stra in strategy_list:
- method_cls_name = stra.pop("method_cls") if "method_cls" in stra
else "default"
+ ), "ToolStrategy should be given as list of dict"
+ for strategy in strategy_list:
+ meta_strategy = msc_utils.copy_dict(strategy)
+ method_cls_name = strategy.pop("method_cls") if "method_cls" in
strategy else "default"
method_cls = msc_utils.get_registered_tool_method(
self.framework(), self.tool_type(), method_cls_name
)
- method_name = stra.pop("method") if "method" in stra else "default"
+ method_name = strategy.pop("method") if "method" in strategy else
"default"
+ method = None
if hasattr(method_cls, method_name):
method = getattr(method_cls, method_name)
- else:
+ if not method:
default_cls = msc_utils.get_registered_tool_method(
MSCFramework.MSC, self.tool_type(), method_cls_name
)
- assert hasattr(
- default_cls, method_name
- ), "Can not find method {} from neighter {} nor {}".format(
- method_name, method_cls, default_cls
- )
- method = getattr(default_cls, method_name)
- tensor_types = stra.pop("tensor_types") if "tensor_types" in stra
else ["all"]
- if "op_types" in stra:
- op_types = stra.pop("op_types")
+ if hasattr(default_cls, method_name):
+ method = getattr(default_cls, method_name)
+ if not method:
+ method = msc_utils.get_registered_func(method_name)
+ assert method, "Can not find method with " + str(method_name)
+ tensor_types = strategy.pop("tensor_types") if "tensor_types" in
strategy else ["all"]
+ if "op_types" in strategy:
+ op_types = strategy.pop("op_types")
marks = [("{}.{}".format(s, t), t) for s, t in
product(op_types, tensor_types)]
- elif "op_names" in stra:
- op_names = stra.pop("op_names")
+ elif "op_names" in strategy:
+ op_names = strategy.pop("op_names")
marks = [("{}.{}".format(s, t), t) for s, t in
product(op_names, tensor_types)]
+ elif "tensor_names" in strategy:
+ tensor_names = strategy.pop("tensor_names")
+ marks = [(n, "all") for n in tensor_names]
else:
marks = [("default", "all")]
- stages = stra.pop("stages") if "stages" in stra else ["default"]
+ stages = strategy.pop("stages") if "stages" in strategy else
["default"]
for mark, t_type in marks:
if mark not in strategys:
- strategys[mark] = Strategy(mark, t_type, self._stage)
+ strategys[mark] = ToolStrategy(mark, t_type, self._stage,
meta_strategy)
for stage in stages:
strategys[mark].add_executor(
- stage, Executor(method_name, method,
copy.deepcopy(stra))
+ stage, ToolExecutor(method_name, method,
copy.deepcopy(strategy))
)
return strategys
@@ -436,41 +445,22 @@ class BaseTool(object):
cache_info = {}
if self.tool_type() in cache_info:
self.load_cache(cache_dir, cache_info[self.tool_type()])
- else:
- graphs, weights = self.load_graphs(graphs, weights)
- self._graphs, self._weights = graphs, {}
+ self._graphs, weights = self._reset(graphs, weights)
+ self._weights = {}
for sub_weights in weights:
self._weights.update(sub_weights)
self._logger.debug(
- "%s load %d graphs and %d weights",
+ "%s reset %d graphs, %d weights",
self.tool_type(),
len(self._graphs),
len(self._weights),
)
- self._reset()
return self._graphs, weights
- def _reset(self):
- """Extra reset for tool"""
-
- return None
-
- def change_stage(self, stage: str):
- """Change the stage of tools and strategy"""
-
- self._stage = stage
- for strategy in self._strategys.values():
- strategy.change_stage(stage)
-
- def destory(self):
- """Destory tool"""
-
- self._graphs, self._weights = [], {}
-
- def load_graphs(
+ def _reset(
self, graphs: List[MSCGraph], weights: List[Dict[str, tvm.nd.array]]
) -> Tuple[List[MSCGraph], List[Dict[str, tvm.nd.array]]]:
- """Load the graphs and weights
+ """Reset the tool
Parameters
----------
@@ -489,6 +479,23 @@ class BaseTool(object):
return graphs, weights
+ def change_stage(self, stage: str):
+ """Change the stage of tool and strategy"""
+
+ self._stage = stage
+ for strategy in self._strategys.values():
+ strategy.change_stage(stage)
+
+ def change_logger(self, logger: logging.Logger):
+ """Change the logger of tool"""
+
+ self._logger = logger
+
+ def destory(self):
+ """Destory tool"""
+
+ self._graphs, self._weights = [], {}
+
def load_cache(self, cache_dir: msc_utils.MSCDirectory, cache_info: dict):
"""Save runner to cache
@@ -677,6 +684,8 @@ class BaseTool(object):
The processed tensor.
"""
+ if not self._enabled:
+ return tensor
if not self._support_scope(scope):
return tensor
strategys = self._get_tensor_strategys(name, consumer)
@@ -780,7 +789,7 @@ class BaseTool(object):
return len(strategys) > 0
def _process_tensor(
- self, tensor: Any, name: str, consumer: str, scope: str, strategys:
List[Strategy]
+ self, tensor: Any, name: str, consumer: str, scope: str, strategys:
List[ToolStrategy]
) -> Any:
"""Process tensor
@@ -794,7 +803,7 @@ class BaseTool(object):
The name of the consumer.
scope: str
The scope mark teacher| student| null.
- strategys: list<Strategy>
+ strategys: list<ToolStrategy>
The strategys for the tensor.
Returns
@@ -832,8 +841,8 @@ class BaseTool(object):
return None
- def update_plan(self, plan: dict):
- """Update the plan
+ def set_plan(self, plan: dict):
+ """Set the plan
Parameters
----------
@@ -841,23 +850,10 @@ class BaseTool(object):
The new plan.
"""
- self._plan.update(plan)
-
- def get_plan(self, name: str) -> dict:
- """Get the plan for name
-
- Parameters
- ----------
- name: str
- The plan name.
-
- Returns
- -------
- plan: dict
- The plan of the name.
- """
-
- return self._plan.get(name, {})
+ if self._plan:
+ self._plan = msc_utils.update_dict(self._plan, plan)
+ else:
+ self._plan = plan
def finalize(self) -> dict:
"""Get the plan"""
@@ -1177,7 +1173,7 @@ class BaseTool(object):
return None
return self._tensor_cache[tensor_id].get(key)
- def _get_tensor_strategys(self, name: str, consumer: str) ->
List[Strategy]:
+ def _get_tensor_strategys(self, name: str, consumer: str) ->
List[ToolStrategy]:
"""Get the strategys by name and consumer
Parameters
@@ -1189,7 +1185,7 @@ class BaseTool(object):
Returns
-------
- strategys: list<Strategy>
+ strategys: list<ToolStrategy>
The strategys for the tensor.
"""
@@ -1222,16 +1218,20 @@ class BaseTool(object):
consumer.optype + ".all",
]
strategys = []
- for n in name_refs:
- if n in self._strategys and
self._strategys[n].support_stage(self._stage):
- strategys.append(self._strategys[n])
+ tensor_strategy = self._strategys.get(tensor_id)
+ if tensor_strategy and tensor_strategy.support_stage(self._stage):
+ strategys.append(tensor_strategy)
+ if not strategys:
+ for n in name_refs:
+ if n in self._strategys and
self._strategys[n].support_stage(self._stage):
+ strategys.append(self._strategys[n])
d_strategy = self._strategys.get("default")
if not strategys and d_strategy and
d_strategy.support_stage(self._stage):
strategys.append(d_strategy)
self._save_tensor_cache(name, consumer, mark, strategys)
return self._get_tensor_cache(name, consumer, mark)
- def _get_tensor_strategy(self, name: str, consumer: str) -> Strategy:
+ def _get_tensor_strategy(self, name: str, consumer: str) -> ToolStrategy:
"""Get the unique strategy by name and consumer
Parameters
@@ -1243,7 +1243,7 @@ class BaseTool(object):
Returns
-------
- strategy: Strategy
+ strategy: ToolStrategy
The unique strategy for the tensor.
"""
@@ -1274,26 +1274,22 @@ class BaseTool(object):
class WeightTool(BaseTool):
"""Basic tool with weight graphs"""
- def _reset(self):
- """Extra reset for tool"""
+ def setup(self) -> dict:
+ """Setup the tool
- super()._reset()
- assert len(self._graphs) == len(
- self._weight_graphs
- ), "Graphs {} mismatch with weight graphs {}".format(
- len(self._graphs), len(self._weight_graphs)
- )
- self._logger.debug("%s load %d weight graphs", self.tool_type(),
len(self._weight_graphs))
- if self.on_debug(2, in_forward=False):
- for idx, graph in enumerate(self._weight_graphs):
- self._logger.debug(
- msc_utils.msg_block("WEIGHT_GRAPH[{}].INFO".format(idx),
graph.inspect())
- )
+ Returns
+ -------
+ info: dict
+ The setup info.
+ """
- def load_graphs(
+ self._weight_graphs = []
+ return super().setup()
+
+ def _reset(
self, graphs: List[MSCGraph], weights: List[Dict[str, tvm.nd.array]]
) -> Tuple[List[MSCGraph], List[Dict[str, tvm.nd.array]]]:
- """Load the graphs and weights
+ """Reset the tool
Parameters
----------
@@ -1312,12 +1308,28 @@ class WeightTool(BaseTool):
The weights
"""
- graphs, weights = super().load_graphs(graphs, weights)
- main_wtypes, relation_wtypes = self._get_wtypes()
- assert main_wtypes, "main_wtypes should be given to build weight
graphs"
- self._weight_graphs = [
- _ffi_api.WeightGraph(graph, main_wtypes, relation_wtypes) for
graph in graphs
- ]
+ graphs, weights = super()._reset(graphs, weights)
+ self._main_wtypes, self._relation_wtypes = self._get_wtypes()
+ assert self._main_wtypes, "main_wtypes should be given to build weight
graphs"
+ if self._weight_graphs:
+ assert len(graphs) == len(
+ self._weight_graphs
+ ), "Graphs {} mismatch with weight graphs {}".format(
+ len(graphs), len(self._weight_graphs)
+ )
+ else:
+ self._weight_graphs = [
+ _ffi_api.WeightGraph(graph, self._main_wtypes,
self._relation_wtypes)
+ for graph in graphs
+ ]
+ self._logger.debug(
+ "%s reset %d weight graphs", self.tool_type(),
len(self._weight_graphs)
+ )
+ if self.on_debug(2, in_forward=False):
+ for idx, graph in enumerate(self._weight_graphs):
+ self._logger.debug(
+ msc_utils.msg_block("WEIGHT_GRAPH[{}].INFO".format(idx),
graph.inspect())
+ )
return graphs, weights
def _get_wtypes(self) -> Tuple[Dict[str, List[str]], Dict[str, str]]:
@@ -1350,6 +1362,12 @@ class WeightTool(BaseTool):
self._weight_graphs = [
WeightGraph.from_json(cache_dir.relpath(f)) for f in
cache_info["weight_graphs"]
]
+ self._logger.debug(
+ "%s load %d weight graphs from %s",
+ self.tool_type(),
+ len(self._weight_graphs),
+ cache_dir,
+ )
def save_cache(self, cache_dir: msc_utils.MSCDirectory) -> dict:
"""Save runner to cache
diff --git a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
b/python/tvm/contrib/msc/core/tools/track/__init__.py
similarity index 90%
copy from python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
copy to python/tvm/contrib/msc/core/tools/track/__init__.py
index 94e19b7e07..2c82a6d486 100644
--- a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
+++ b/python/tvm/contrib/msc/core/tools/track/__init__.py
@@ -14,6 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""tvm.contrib.msc.framework.tensorflow.tools"""
+"""tvm.contrib.msc.core.tools.track"""
-from .prune import *
+from .tracker import *
+from .method import *
diff --git a/python/tvm/contrib/msc/core/tools/track/method.py
b/python/tvm/contrib/msc/core/tools/track/method.py
new file mode 100644
index 0000000000..aaa07b3812
--- /dev/null
+++ b/python/tvm/contrib/msc/core/tools/track/method.py
@@ -0,0 +1,102 @@
+# 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.
+# pylint: disable=unused-argument
+"""tvm.contrib.msc.core.tools.track.method"""
+
+from typing import List, Dict
+import numpy as np
+
+from tvm.contrib.msc.core.tools.tool import ToolType, BaseTool
+from tvm.contrib.msc.core.utils.namespace import MSCFramework
+from tvm.contrib.msc.core import utils as msc_utils
+
+
+class TrackMethod(object):
+ """Default track method"""
+
+ @classmethod
+ def save_compared(
+ cls,
+ tracker: BaseTool,
+ data: np.ndarray,
+ name: str,
+ consumer: str,
+ compare_to: Dict[str, List[str]],
+ ) -> np.ndarray:
+ """Compare and save the data
+
+ Parameters
+ ----------
+ tracker: BaseTracker
+ The tracker
+ data: np.ndarray
+ The source data.
+ name: str
+ The name of the tensor.
+ consumer: str
+ The name of the consumer.
+ stage: str
+ The current stage of tool.
+ compare_to: dict
+ The compare config
+ dataset: MSCDirectory
+ The root dir
+
+ Returns
+ -------
+ plan: dict
+ The plan of the tensor.
+ """
+
+ data = msc_utils.cast_array(data)
+ config = {"info": msc_utils.inspect_array(data)}
+ # save the data
+ tracker._saver.save_datas({name: data}, tracker._forward_cnt)
+ tracker.debug_tensor(data, name, consumer, "save")
+ # compare datas
+ if tracker._stage in compare_to:
+ diffs = {}
+ for stage in compare_to[tracker._stage]:
+ if stage in tracker._loaders:
+ if not tracker._loaders[stage].has_data(name,
tracker._forward_cnt):
+ continue
+ golden = tracker._loaders[stage].load_data(name,
tracker._forward_cnt)
+ report = msc_utils.compare_arrays({name: golden}, {name:
data})
+ diff_msg = "{}diff to {} -> {}".format(
+ tracker.msg_mark(), stage, report["info"][name]
+ )
+ if report["passed"] == 0:
+ tracker._logger.info(diff_msg)
+ elif tracker.on_debug(3):
+ tracker._logger.debug(diff_msg)
+ diffs[stage] = {
+ "pass": report["passed"] == 1,
+ "info": msc_utils.inspect_array(np.abs(golden - data)),
+ }
+ config["diffs"] = diffs
+ return config
+
+ @classmethod
+ def framework(cls):
+ return MSCFramework.MSC
+
+ @classmethod
+ def tool_type(cls):
+ return ToolType.TRACKER
+
+
+msc_utils.register_tool_method(TrackMethod)
diff --git a/python/tvm/contrib/msc/core/tools/track/tracker.py
b/python/tvm/contrib/msc/core/tools/track/tracker.py
new file mode 100644
index 0000000000..442ac6f508
--- /dev/null
+++ b/python/tvm/contrib/msc/core/tools/track/tracker.py
@@ -0,0 +1,185 @@
+# 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.
+"""tvm.contrib.msc.core.tools.track.tracker"""
+
+from typing import Any, List
+from tvm.contrib.msc.core.tools.tool import ToolType, BaseTool, ToolStrategy
+from tvm.contrib.msc.core import utils as msc_utils
+
+
+class BaseTracker(BaseTool):
+ """Base tracker for all"""
+
+ def setup(self) -> dict:
+ """Setup the tool
+
+ Returns
+ -------
+ info: dict
+ The setup info.
+ """
+
+ data_folder = msc_utils.get_dataset_dir().create_dir("Track")
+ self._loaders = {}
+ for folder in data_folder.listdir():
+ if msc_utils.is_simple_dataset(data_folder.relpath(folder)):
+ self._loaders[folder] =
msc_utils.SimpleDataLoader(data_folder.relpath(folder))
+ self._saver =
msc_utils.SimpleDataSaver(data_folder.relpath(self._stage))
+ self._max_iter = self._options.get("max_iter", 1)
+ info = super().setup()
+ info.update({"saver": self._saver, "loaders": self._loaders})
+ return info
+
+ def finalize(self) -> dict:
+ """Get the plan"""
+
+ self._saver.finalize()
+ return super().finalize()
+
+ def _execute_after_forward(self, output: Any) -> Any:
+ """Execute after model forward
+
+ Parameters
+ ----------
+ output: Any
+ The output reference of the model.
+
+ Returns
+ -------
+ output: Any
+ The modified output reference.
+ """
+
+ if self._forward_cnt < self._max_iter:
+ passed = {}
+ for info in self._plan.values():
+ if "diffs" not in info[self._stage]:
+ continue
+ for stage, p_info in info[self._stage]["diffs"].items():
+ if stage not in passed:
+ passed[stage] = {"total": 0, "passed": 0}
+ passed[stage]["total"] += 1
+ if p_info["pass"]:
+ passed[stage]["passed"] += 1
+ msg = "Track({})[{}] {} datas".format(self._stage,
self._forward_cnt, len(self._plan))
+ if passed:
+ msg += ", passed -> "
+ msg += "; ".join(
+ ["{}: {}/{}".format(s, i["passed"], i["total"]) for s, i
in passed.items()]
+ )
+ self._logger.info(msg)
+ return output
+
+ def _check_tensor(self, name: str, consumer: str) -> bool:
+ """Check if the tensor should be processed
+
+ Parameters
+ -------
+ name: str
+ The name of the tensor.
+ consumer: str
+ The name of the consumer.
+
+ Returns
+ -------
+ vaild: bool
+ Whether to process the tensor.
+ """
+
+ if self._forward_cnt >= self._max_iter:
+ return False
+ strategy = self._get_tensor_strategy(name, consumer)
+ if not strategy:
+ return False
+ compare_to = strategy.get_config().get("compare_to", {})
+ if self._stage in compare_to:
+ return True
+ for stages in compare_to.values():
+ if self._stage in stages:
+ return True
+ return False
+
+ def _process_tensor(
+ self, tensor: Any, name: str, consumer: str, scope: str, strategys:
List[ToolStrategy]
+ ) -> Any:
+ """Process tensor
+
+ Parameters
+ -------
+ tensor: Any
+ Tensor in framework
+ name: str
+ The name of the tensor.
+ consumer: str
+ The name of the consumer.
+ scope: str
+ The scope mark teacher| student| null.
+ strategys: list<ToolStrategy>
+ The strategys for the tensor.
+
+ Returns
+ -------
+ tensor: Any
+ The processed tensor.
+ """
+
+ return self._track_tensor(tensor, name, consumer, strategys)
+
+ def _track_tensor(
+ self, tensor: Any, name: str, consumer: str, strategys:
List[ToolStrategy]
+ ) -> Any:
+ """Process tensor
+
+ Parameters
+ -------
+ tensor: Any
+ Tensor in framework
+ name: str
+ The name of the tensor.
+ consumer: str
+ The name of the consumer.
+ strategys: list<ToolStrategy>
+ The strategys for the tensor.
+
+ Returns
+ -------
+ tensor: Any
+ The processed tensor.
+ """
+
+ if self._stage in self._plan.get(name, {}):
+ return tensor
+ if name not in self._plan:
+ self._plan[name] = {}
+ plan = {}
+ for strategy in strategys:
+ plan.update(strategy(self, tensor, name, consumer))
+ self._plan[name][self._stage] = plan
+ return tensor
+
+ @classmethod
+ def tool_type(cls):
+ return ToolType.TRACKER
+
+
+class DefaultTracker(BaseTracker):
+ @classmethod
+ def tool_style(cls):
+ return "default"
+
+
+msc_utils.register_tool_cls(DefaultTracker)
diff --git a/python/tvm/contrib/msc/core/utils/info.py
b/python/tvm/contrib/msc/core/utils/info.py
index 782be26049..440789f856 100644
--- a/python/tvm/contrib/msc/core/utils/info.py
+++ b/python/tvm/contrib/msc/core/utils/info.py
@@ -238,6 +238,40 @@ def load_dict(str_dict: str, flavor: str = "json") -> dict:
return dict_obj
+def update_dict(
+ src_dict: dict, new_dict: dict, recursive: bool = True, soft_update: bool
= True
+) -> dict:
+ """Update src_dict with new_dict.
+
+ Parameters
+ ----------
+ src_dict: dict
+ The source dict.
+ new_dict: dict
+ The new dict.
+ recursive: bool
+ Whether to update the dict recursive.
+ soft_update: bool
+ Whether to update the source dict, False to force update.
+
+ Returns
+ -------
+ dict_obj: dict
+ The updated dict.
+ """
+
+ assert isinstance(src_dict, dict) and isinstance(
+ new_dict, dict
+ ), "update_dict only support dict, get src {} and new
{}".format(type(src_dict), type(new_dict))
+ for k, v in new_dict.items():
+ if isinstance(v, dict):
+ v = update_dict(src_dict.get(k, {}), v, recursive, soft_update)
+ src_dict[k] = v
+ elif not soft_update or k not in src_dict:
+ src_dict[k] = v
+ return src_dict
+
+
def dump_dict(dict_obj: dict, flavor: str = "dmlc") -> str:
"""Dump the config to string.
@@ -288,7 +322,8 @@ def dump_dict(dict_obj: dict, flavor: str = "dmlc") -> str:
lines.append("{}{}: {}".format(indent * " ", k, v))
return lines
- return "\n".join(_get_lines(dict_obj))
+ lines = _get_lines(dict_obj) or [" {}: {}".format(k, v) for k, v in
dict_obj.items()]
+ return "\n".join(lines)
return json.dumps(dict_obj)
diff --git a/python/tvm/contrib/msc/core/utils/log.py
b/python/tvm/contrib/msc/core/utils/log.py
index 525f1706f9..a406db806a 100644
--- a/python/tvm/contrib/msc/core/utils/log.py
+++ b/python/tvm/contrib/msc/core/utils/log.py
@@ -52,7 +52,7 @@ class IOLogger(object):
raise Exception(msg)
-def create_file_logger(level=logging.INFO, path: str = None) -> logging.Logger:
+def create_file_logger(level: Union[str, int] = logging.INFO, path: str =
None) -> logging.Logger:
"""Create file logger
Parameters
@@ -68,6 +68,16 @@ def create_file_logger(level=logging.INFO, path: str = None)
-> logging.Logger:
The logger.
"""
+ if isinstance(level, str):
+ if level == "debug":
+ level = logging.DEBUG
+ elif level == "info":
+ level = logging.INFO
+ elif level == "warn":
+ level = logging.WARN
+ else:
+ raise Exception("Unexcept verbose {}, should be debug| info| warn")
+
path = path or os.path.join(get_workspace(), "MSC_LOG")
log_name = os.path.basename(path)
logger = logging.getLogger(log_name)
@@ -104,15 +114,6 @@ def set_global_logger(level: Union[str, int] =
logging.INFO, path: str = None) -
The logger.
"""
- if isinstance(level, str):
- if level == "debug":
- level = logging.DEBUG
- elif level == "info":
- level = logging.INFO
- elif level == "warn":
- level = logging.WARN
- else:
- raise Exception("Unexcept verbose {}, should be debug| info| warn")
logger = create_file_logger(level, path)
MSCMap.set(MSCKey.GLOBALE_LOGGER, logger)
return logger
diff --git a/python/tvm/contrib/msc/core/utils/message.py
b/python/tvm/contrib/msc/core/utils/message.py
index 69c31c807e..4f93d402a0 100644
--- a/python/tvm/contrib/msc/core/utils/message.py
+++ b/python/tvm/contrib/msc/core/utils/message.py
@@ -124,6 +124,29 @@ def get_duration() -> dict:
return duration
+def msg_table(title: str, msg: str, width: int = 100):
+ """Log message in table format
+
+ Parameters
+ ----------
+ title: str
+ The title of the block
+ msg: str
+ The message to log.
+ width: int
+ The max width of block message
+
+ Returns
+ -------
+ msg: str
+ The block message.
+ """
+
+ if isinstance(msg, dict):
+ msg = dump_dict(msg, "table:" + str(width))
+ return "\n{0} {1} {0}\n{2}\n".format("-" * 20, title.center(40), msg)
+
+
def msg_block(title: str, msg: str, width: int = 100):
"""Log message in block format
diff --git a/python/tvm/contrib/msc/framework/tensorflow/runtime/runner.py
b/python/tvm/contrib/msc/framework/tensorflow/runtime/runner.py
index c686647bfe..4617c5d351 100644
--- a/python/tvm/contrib/msc/framework/tensorflow/runtime/runner.py
+++ b/python/tvm/contrib/msc/framework/tensorflow/runtime/runner.py
@@ -93,7 +93,7 @@ class TensorflowRunner(ModelRunner):
-------
graphs: list<MSCgraph>
The msc graphs.
- weights: list<dic<str, tvm.nd.array>>
+ weights: list<dict<str, tvm.nd.array>>
The weights
Returns
diff --git a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
b/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
index 94e19b7e07..d25cfd4e67 100644
--- a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
+++ b/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
@@ -17,3 +17,4 @@
"""tvm.contrib.msc.framework.tensorflow.tools"""
from .prune import *
+from .track import *
diff --git a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
b/python/tvm/contrib/msc/framework/tensorflow/tools/track/__init__.py
similarity index 90%
copy from python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
copy to python/tvm/contrib/msc/framework/tensorflow/tools/track/__init__.py
index 94e19b7e07..e8787fb666 100644
--- a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
+++ b/python/tvm/contrib/msc/framework/tensorflow/tools/track/__init__.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""tvm.contrib.msc.framework.tensorflow.tools"""
+"""tvm.contrib.msc.framework.tensorflow.tools.track"""
-from .prune import *
+from .tracker import *
diff --git a/python/tvm/contrib/msc/framework/tensorflow/tools/track/tracker.py
b/python/tvm/contrib/msc/framework/tensorflow/tools/track/tracker.py
new file mode 100644
index 0000000000..7023322681
--- /dev/null
+++ b/python/tvm/contrib/msc/framework/tensorflow/tools/track/tracker.py
@@ -0,0 +1,55 @@
+# 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.
+"""tvm.contrib.msc.framework.tensorflow.tools.track.tracker"""
+
+from tvm.contrib.msc.core.tools.tool import ToolType
+from tvm.contrib.msc.core.tools.track import BaseTracker
+from tvm.contrib.msc.core.utils.namespace import MSCFramework
+from tvm.contrib.msc.core import utils as msc_utils
+
+
+class TensorflowTrackerFactory(object):
+ """Tracker factory for tensorflow"""
+
+ def create(self, base_cls: BaseTracker) -> BaseTracker:
+ """Create adaptive tracker
+
+ Parameters
+ ----------
+ base_cls: BaseTracker
+ The base tracker class
+
+ Returns
+ -------
+ tracker_cls: BaseTracker
+ The tracker class.
+ """
+
+ class Tracker(base_cls):
+ """Adaptive tracker for tensorflow"""
+
+ @classmethod
+ def framework(cls):
+ return MSCFramework.TENSORFLOW
+
+ return Tracker
+
+
+factory = TensorflowTrackerFactory()
+tools = msc_utils.get_registered_tool_cls(MSCFramework.MSC, ToolType.TRACKER,
tool_style="all")
+for tool in tools.values():
+ msc_utils.register_tool_cls(factory.create(tool))
diff --git a/python/tvm/contrib/msc/framework/tensorrt/tools/__init__.py
b/python/tvm/contrib/msc/framework/tensorrt/tools/__init__.py
index 0247da2642..ecc82bc40f 100644
--- a/python/tvm/contrib/msc/framework/tensorrt/tools/__init__.py
+++ b/python/tvm/contrib/msc/framework/tensorrt/tools/__init__.py
@@ -17,3 +17,4 @@
"""tvm.contrib.msc.framework.tensorrt.tools"""
from .prune import *
+from .track import *
diff --git a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
b/python/tvm/contrib/msc/framework/tensorrt/tools/track/__init__.py
similarity index 91%
copy from python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
copy to python/tvm/contrib/msc/framework/tensorrt/tools/track/__init__.py
index 94e19b7e07..88897c1da5 100644
--- a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
+++ b/python/tvm/contrib/msc/framework/tensorrt/tools/track/__init__.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""tvm.contrib.msc.framework.tensorflow.tools"""
+"""tvm.contrib.msc.framework.tensorrt.tools.track"""
-from .prune import *
+from .tracker import *
diff --git a/python/tvm/contrib/msc/framework/tensorrt/tools/track/tracker.py
b/python/tvm/contrib/msc/framework/tensorrt/tools/track/tracker.py
new file mode 100644
index 0000000000..10ae794ca0
--- /dev/null
+++ b/python/tvm/contrib/msc/framework/tensorrt/tools/track/tracker.py
@@ -0,0 +1,159 @@
+# 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.
+# pylint: disable=unused-argument
+"""tvm.contrib.msc.framework.tensorrt.tools.track.tracker"""
+
+from typing import Dict, List
+
+from tvm.contrib.msc.core.tools.tool import ToolType, ToolStrategy
+from tvm.contrib.msc.core.tools.track import BaseTracker
+from tvm.contrib.msc.core.utils.namespace import MSCFramework
+from tvm.contrib.msc.core import utils as msc_utils
+
+
+class TensorRTTrackerFactory(object):
+ """Tracker factory for tensorrt"""
+
+ def create(self, base_cls: BaseTracker) -> BaseTracker:
+ """Create adaptive tracker
+
+ Parameters
+ ----------
+ base_cls: BaseTracker
+ The base tracker class
+
+ Returns
+ -------
+ tracker_cls: BaseTracker
+ The tracker class.
+ """
+
+ class Tracker(base_cls):
+ """Adaptive tracker for tensorrt"""
+
+ def _execute_before_build(self, codegen_context: dict) -> dict:
+ """Execute before model build
+
+ Parameters
+ ----------
+ codegen_context: dict
+ The context.
+
+ Returns
+ ----------
+ codegen_context: dict
+ The processed context.
+ """
+
+ self._track_tensors = {}
+ super()._execute_before_build(codegen_context)
+
+ def _execute_before_forward(self, step_context: dict) -> dict:
+ """Execute before model forward
+
+ Parameters
+ ----------
+ step_context: dict
+ The context.
+
+ Returns
+ ----------
+ step_context: dict
+ The processed context.
+ """
+
+ for name, data in step_context["datas"].items():
+ if name not in self._track_tensors:
+ continue
+ consumer = self._track_tensors[name]["consumer"]
+ strategys = self._get_tensor_strategys(name, consumer)
+ self._track_tensor(data.asnumpy(), name, consumer,
strategys)
+ return super()._execute_before_forward(step_context)
+
+ def _execute_after_forward(self, step_context: dict) -> dict:
+ """Execute after model forward
+
+ Parameters
+ ----------
+ step_context: dict
+ The context.
+
+ Returns
+ ----------
+ step_context: dict
+ The processed context.
+ """
+
+ for name, data in step_context["datas"].items():
+ if name not in self._track_tensors:
+ continue
+ consumer = self._track_tensors[name]["consumer"]
+ strategys = self._get_tensor_strategys(name, consumer)
+ self._track_tensor(data.asnumpy(), name, consumer,
strategys)
+ return super()._execute_after_forward(step_context)
+
+ def _process_tensor(
+ self,
+ tensor_ctx: Dict[str, str],
+ name: str,
+ consumer: str,
+ scope: str,
+ strategys: List[ToolStrategy],
+ ) -> Dict[str, str]:
+ """Process tensor
+
+ Parameters
+ -------
+ tensor_ctx: dict<str, str>
+ Tensor describe items.
+ name: str
+ The name of the tensor.
+ consumer: str
+ The name of the consumer.
+ scope: str
+ The scope mark teacher| student| null.
+ strategys: list<ToolStrategy>
+ The strategys for the tensor.
+
+ Returns
+ -------
+ tensor_ctx: dict<str, str>
+ Tensor items with processed.
+ """
+
+ if self.is_weight(name):
+ return self._track_tensor(self.get_data(name), name,
consumer, strategys)
+ if name not in self._track_tensors:
+ self._track_tensors[name] = {
+ "consumer": consumer,
+ }
+ tensor_ctx["processed"].append(
+ "{}->markOutput(*{});".format(tensor_ctx["ctx"],
tensor_ctx["tensor"])
+ )
+ return tensor_ctx
+
+ @classmethod
+ def framework(cls):
+ return MSCFramework.TENSORRT
+
+ return Tracker
+
+
+factory = TensorRTTrackerFactory()
+tools = msc_utils.get_registered_tool_cls(MSCFramework.MSC, ToolType.TRACKER,
tool_style="all")
+for tool in tools.values():
+ msc_utils.register_tool_cls(factory.create(tool))
diff --git a/python/tvm/contrib/msc/framework/torch/frontend/translate.py
b/python/tvm/contrib/msc/framework/torch/frontend/translate.py
index 2dce394708..3ac1b81a2c 100644
--- a/python/tvm/contrib/msc/framework/torch/frontend/translate.py
+++ b/python/tvm/contrib/msc/framework/torch/frontend/translate.py
@@ -57,7 +57,7 @@ def set_weight_alias(graph: MSCGraph) -> MSCGraph:
alias = node.name.replace(".", "_") + ".running_var"
else:
alias = node.name.replace(".", "_") + "." + ref
- weight.set_alias(alias)
+ graph.set_tensor_alias(weight, alias)
return graph
diff --git a/python/tvm/contrib/msc/framework/torch/tools/__init__.py
b/python/tvm/contrib/msc/framework/torch/tools/__init__.py
index ff26491f54..dda1e13822 100644
--- a/python/tvm/contrib/msc/framework/torch/tools/__init__.py
+++ b/python/tvm/contrib/msc/framework/torch/tools/__init__.py
@@ -17,3 +17,4 @@
"""tvm.contrib.msc.framework.torch.tools"""
from .prune import *
+from .track import *
diff --git a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
b/python/tvm/contrib/msc/framework/torch/tools/track/__init__.py
similarity index 91%
copy from python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
copy to python/tvm/contrib/msc/framework/torch/tools/track/__init__.py
index 94e19b7e07..55951fe3a9 100644
--- a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
+++ b/python/tvm/contrib/msc/framework/torch/tools/track/__init__.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""tvm.contrib.msc.framework.tensorflow.tools"""
+"""tvm.contrib.msc.framework.torch.tools.track"""
-from .prune import *
+from .tracker import *
diff --git a/python/tvm/contrib/msc/framework/torch/tools/track/tracker.py
b/python/tvm/contrib/msc/framework/torch/tools/track/tracker.py
new file mode 100644
index 0000000000..0fa065153b
--- /dev/null
+++ b/python/tvm/contrib/msc/framework/torch/tools/track/tracker.py
@@ -0,0 +1,55 @@
+# 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.
+"""tvm.contrib.msc.framework.torch.tools.track.tracker"""
+
+from tvm.contrib.msc.core.tools.tool import ToolType
+from tvm.contrib.msc.core.tools.track import BaseTracker
+from tvm.contrib.msc.core.utils.namespace import MSCFramework
+from tvm.contrib.msc.core import utils as msc_utils
+
+
+class TorchTrackerFactory(object):
+ """Tracker factory for torch"""
+
+ def create(self, base_cls: BaseTracker) -> BaseTracker:
+ """Create adaptive tracker
+
+ Parameters
+ ----------
+ base_cls: BaseTracker
+ The base tracker class
+
+ Returns
+ -------
+ tracker_cls: BaseTracker
+ The tracker class.
+ """
+
+ class Tracker(base_cls):
+ """Adaptive tracker for torch"""
+
+ @classmethod
+ def framework(cls):
+ return MSCFramework.TORCH
+
+ return Tracker
+
+
+factory = TorchTrackerFactory()
+tools = msc_utils.get_registered_tool_cls(MSCFramework.MSC, ToolType.TRACKER,
tool_style="all")
+for tool in tools.values():
+ msc_utils.register_tool_cls(factory.create(tool))
diff --git a/python/tvm/contrib/msc/framework/tvm/tools/__init__.py
b/python/tvm/contrib/msc/framework/tvm/tools/__init__.py
index 91f07fd581..226ae3102d 100644
--- a/python/tvm/contrib/msc/framework/tvm/tools/__init__.py
+++ b/python/tvm/contrib/msc/framework/tvm/tools/__init__.py
@@ -17,3 +17,4 @@
"""tvm.contrib.msc.framework.tvm.tools"""
from .prune import *
+from .track import *
diff --git a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
b/python/tvm/contrib/msc/framework/tvm/tools/track/__init__.py
similarity index 91%
copy from python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
copy to python/tvm/contrib/msc/framework/tvm/tools/track/__init__.py
index 94e19b7e07..f99f09767d 100644
--- a/python/tvm/contrib/msc/framework/tensorflow/tools/__init__.py
+++ b/python/tvm/contrib/msc/framework/tvm/tools/track/__init__.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""tvm.contrib.msc.framework.tensorflow.tools"""
+"""tvm.contrib.msc.framework.tvm.tools.track"""
-from .prune import *
+from .tracker import *
diff --git a/python/tvm/contrib/msc/framework/tvm/tools/track/tracker.py
b/python/tvm/contrib/msc/framework/tvm/tools/track/tracker.py
new file mode 100644
index 0000000000..cf5ab49e82
--- /dev/null
+++ b/python/tvm/contrib/msc/framework/tvm/tools/track/tracker.py
@@ -0,0 +1,155 @@
+# 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.
+# pylint: disable=unused-argument
+"""tvm.contrib.msc.framework.tvm.tools.track.tracker"""
+
+from typing import List, Union
+
+import tvm
+from tvm.contrib.msc.core.tools.tool import ToolType, ToolStrategy
+from tvm.contrib.msc.core.tools.track import BaseTracker
+from tvm.contrib.msc.core.utils.namespace import MSCFramework
+from tvm.contrib.msc.core import utils as msc_utils
+
+
+class TVMTrackerFactory(object):
+ """Tracker factory for tvm"""
+
+ def create(self, base_cls: BaseTracker) -> BaseTracker:
+ """Create adaptive tracker
+
+ Parameters
+ ----------
+ base_cls: BaseTracker
+ The base tracker class
+
+ Returns
+ -------
+ tracker_cls: BaseTracker
+ The tracker class.
+ """
+
+ class Tracker(base_cls):
+ """Adaptive tracker for tvm"""
+
+ def _execute_before_build(self, block_builder:
tvm.relax.BlockBuilder):
+ """Execute before model build
+
+ Parameters
+ ----------
+ block_builder: tvm.relax.BlockBuilder
+ The block builder.
+ """
+
+ self._block_builder = block_builder
+ self._track_tensors, self._track_names = {}, []
+ super()._execute_before_build(block_builder)
+
+ def _execute_after_build(
+ self, output: Union[tvm.relax.Var, List[tvm.relax.DataflowVar]]
+ ) -> List[tvm.relax.Var]:
+ """Execute after model build
+
+ Parameters
+ ----------
+ output: var or list<var>
+ The output var of the model.
+
+ Returns
+ -------
+ outputs: list<var>
+ The modified outputs var.
+ """
+
+ self._track_names = list(sorted(self._track_tensors.keys()))
+ track_tensors = [self._track_tensors[o]["tensor"] for o in
self._track_names]
+ if isinstance(output, tvm.relax.Var):
+ return super()._execute_after_build([output] +
track_tensors)
+ return super()._execute_after_build(output + track_tensors)
+
+ def _execute_after_forward(
+ self, outputs: List[tvm.runtime.NDArray]
+ ) -> Union[tvm.runtime.NDArray, List[tvm.runtime.NDArray]]:
+ """Execute after model forward
+
+ Parameters
+ ----------
+ outputs: list<np.ndarray>
+ The output datas.
+
+ Returns
+ -------
+ output: np.ndarray or list<np.ndarray>
+ The modified output ndarray.
+ """
+
+ output_num = len(outputs) - len(self._track_names)
+ for data, name in zip(outputs[output_num:], self._track_names):
+ consumer = self._track_tensors[name]["consumer"]
+ strategys = self._get_tensor_strategys(name, consumer)
+ self._track_tensor(data, name, consumer, strategys)
+ if output_num == 1:
+ return super()._execute_after_forward(outputs[0])
+ return super()._execute_after_forward(outputs[:output_num])
+
+ def _process_tensor(
+ self,
+ tensor: tvm.relax.DataflowVar,
+ name: str,
+ consumer: str,
+ scope: str,
+ strategys: List[ToolStrategy],
+ ) -> tvm.relax.DataflowVar:
+ """Process tensor
+
+ Parameters
+ -------
+ tensor: Any
+ Tensor in framework
+ name: str
+ The name of the tensor.
+ consumer: str
+ The name of the consumer.
+ scope: str
+ The scope mark teacher| student| null.
+ strategys: list<ToolStrategy>
+ The strategys for the tensor.
+
+ Returns
+ -------
+ tensor: Any
+ The processed tensor.
+ """
+
+ if self.is_weight(name):
+ return self._track_tensor(self.get_data(name), name,
consumer, strategys)
+ if name not in self._track_tensors:
+ self._track_tensors[name] = {"consumer": consumer,
"tensor": tensor}
+ self._track_names.append(name)
+ return tensor
+
+ @classmethod
+ def framework(cls):
+ return MSCFramework.TVM
+
+ return Tracker
+
+
+factory = TVMTrackerFactory()
+tools = msc_utils.get_registered_tool_cls(MSCFramework.MSC, ToolType.TRACKER,
tool_style="all")
+for tool in tools.values():
+ msc_utils.register_tool_cls(factory.create(tool))
diff --git a/python/tvm/contrib/msc/pipeline/manager.py
b/python/tvm/contrib/msc/pipeline/manager.py
index 4f31eeacfa..bbd6d452ad 100644
--- a/python/tvm/contrib/msc/pipeline/manager.py
+++ b/python/tvm/contrib/msc/pipeline/manager.py
@@ -114,15 +114,6 @@ class BaseManager(object):
config = self._update_runner_config(config, stage)
config = self._update_tool_config(config)
- def _get_tool_stage(tool_type: str) -> str:
- if tool_type == ToolType.PRUNER:
- return MSCStage.PRUNE
- if tool_type == ToolType.QUANTIZER:
- return MSCStage.QUANTIZE
- if tool_type == ToolType.DISTILLER:
- return MSCStage.DISTILL
- return tool_type
-
def _set_debug_level(stage: str, stage_config: dict, default: int =
None) -> dict:
if "debug_level" in stage_config:
debug_levels[stage] = stage_config["debug_level"]
@@ -141,7 +132,7 @@ class BaseManager(object):
if t_type not in config["optimize"]:
continue
debug_levels = _set_debug_level(
- _get_tool_stage(t_type), config["optimize"][t_type],
debug_level
+ self._get_tool_stage(t_type), config["optimize"][t_type],
debug_level
)
ordered_keys = [
"model_type",
@@ -378,16 +369,14 @@ class BaseManager(object):
The runner.
"""
+ runner_cls = self._get_runner_cls(stage_config["run_type"])
+
+ def _tool_enabled(tool_type: str) -> bool:
+ return tool_type in stage_config and
runner_cls.support_tool(tool_type)
+
# run prune
- if ToolType.PRUNER in stage_config:
- self._tools_config[ToolType.PRUNER] = stage_config[ToolType.PRUNER]
- plan_file = stage_config[ToolType.PRUNER]["plan_file"]
- if os.path.isfile(plan_file):
- self._logger.info("Skip %s with plan_file %s",
ToolType.PRUNER, plan_file)
- else:
- msc_utils.time_stamp(MSCStage.PRUNE)
- runner = self._create_tool_runner(MSCStage.PRUNE, stage_config)
- runner.apply_tool(ToolType.PRUNER, self._data_loader)
+ if _tool_enabled(ToolType.PRUNER):
+ self._apply_tool(ToolType.PRUNER, stage_config)
# optimize and get the runner
msc_utils.time_stamp(MSCStage.OPTIMIZE)
@@ -488,7 +477,6 @@ class BaseManager(object):
if self._runner:
self._runner.destory()
- debug_level = self._debug_levels.get(stage, 0)
cache_dir = msc_utils.get_cache_dir().create_dir(stage) if use_cache
else None
tools_config = tools_config or {}
msc_utils.time_stamp(stage + ".build", False)
@@ -496,14 +484,14 @@ class BaseManager(object):
run_config = msc_utils.copy_dict(stage_config.get("run_config"))
if "generate_config" not in run_config:
run_config["generate_config"] = {}
- run_config["generate_config"].update(
- {
- "build_folder": msc_utils.get_build_dir().create_dir(
- stage, cleanup=(debug_level == 0)
- ),
- }
+ cleanup = self._debug_levels.get(stage, 0) == 0
+ run_config["generate_config"]["build_folder"] =
msc_utils.get_build_dir().create_dir(
+ stage, cleanup=cleanup
)
- self._logger.debug("Create runner(%s) by %s(%s)", stage,
runner_cls.__name__, run_config)
+ opt_config = self._config.get("optimize", {})
+ if ToolType.TRACKER in opt_config and
runner_cls.support_tool(ToolType.TRACKER):
+ tools_config = {**tools_config, ToolType.TRACKER:
opt_config[ToolType.TRACKER]}
+ # Build runner
runner = runner_cls(
self._relax_mod,
tools_config=tools_config,
@@ -519,10 +507,12 @@ class BaseManager(object):
self._report["profile"][stage] = self._profile_runner(runner,
stage_config)
if use_cache:
runner.save_cache(cache_dir)
+ if runner.get_tool(ToolType.TRACKER):
+ runner.apply_tool(ToolType.TRACKER)
return runner
- def _create_tool_runner(self, tool_type: str, stage_config: dict) ->
BaseRunner:
- """Create runner with tool.
+ def _apply_tool(self, tool_type: str, stage_config: dict, add_tool: bool =
True) -> str:
+ """Apply tool with runner
Parameters
----------
@@ -530,24 +520,41 @@ class BaseManager(object):
The tool type.
stage_config: dict
The config of this stage.
+ add_tool: bool
+ Whether to add tool in self._tools.
Returns
-------
- runner: BaseRunner
- The runner.
+ plan_file: str
+ The plan_file path.
"""
+ assert tool_type in stage_config, "Can not find config for tool " +
str(tool_type)
+ tool_stage, tool_config = self._get_tool_stage(tool_type),
stage_config[tool_type]
+ plan_file = tool_config["plan_file"]
+ if "gym_configs" in tool_config:
+ gym_configs = tool_config.pop("gym_configs")
+ else:
+ gym_configs = None
+ if add_tool:
+ self._tools_config[tool_type] = tool_config
+ tools_config = self._tools_config
+ else:
+ tools_config = {**self._tools_config, tool_type: tool_config}
+ if os.path.isfile(plan_file):
+ self._logger.info("Skip %s with plan %s", tool_type, plan_file)
+ return plan_file
+ msc_utils.time_stamp(tool_stage)
t_stage_config = {
"run_type": stage_config["run_type"],
"run_config": stage_config["run_config"],
}
- return self._create_runner(
- tool_type,
- t_stage_config,
- tools_config=self._tools_config,
- profile=False,
- use_cache=False,
+ runner = self._create_runner(
+ tool_stage, t_stage_config, tools_config=tools_config,
profile=False, use_cache=False
)
+ if gym_configs:
+ raise NotImplementedError("Gym is not implemented")
+ return runner.apply_tool(tool_type, self._data_loader)
def _profile_runner(self, runner: BaseRunner, stage_config: str) -> dict:
"""Profile the runner.
@@ -590,7 +597,7 @@ class BaseManager(object):
pass_rate = float(passed) / total
report["accuracy"] = "{}/{}({:.2f}%)".format(passed, total,
pass_rate * 100)
title = "Check({}) pass {}".format(stage, report["accuracy"])
- self._logger.debug(msc_utils.msg_block(title, acc_report))
+ self._logger.debug(msc_utils.msg_block(title, acc_report, width=0))
msg += " acc {} iters -> {}".format(len(loader),
report["accuracy"])
if runner.get_tool(ToolType.PRUNER) or
runner.get_tool(ToolType.QUANTIZER):
self._logger.debug("Disable accuracy check(%s) by tools",
stage)
@@ -603,7 +610,12 @@ class BaseManager(object):
)
)
- benchmark_config = profile_config.get("benchmark", {})
+ # benchmark model
+ if runner.get_tool(ToolType.TRACKER):
+ benchmark_config = None
+ self._logger.debug("Disable benchmark(%s) by tools", stage)
+ else:
+ benchmark_config = profile_config.get("benchmark", {})
if benchmark_config:
for _ in range(benchmark_config.get("warm_up", 10)):
runner.run(self._sample_inputs)
@@ -764,6 +776,28 @@ class BaseManager(object):
)
return config
+ def _get_tool_stage(self, tool_type: str) -> str:
+ """Map the stage according to tool_type
+
+ Parameters
+ ----------
+ tool_type: str
+ The tool type.
+
+ Returns
+ -------
+ stage: str
+ The stage.
+ """
+
+ if tool_type == ToolType.PRUNER:
+ return MSCStage.PRUNE
+ if tool_type == ToolType.QUANTIZER:
+ return MSCStage.QUANTIZE
+ if tool_type == ToolType.DISTILLER:
+ return MSCStage.DISTILL
+ return tool_type
+
def get_runnable(self, ret_type: str = "runner") -> Any:
"""Return object by type.
diff --git a/src/contrib/msc/core/ir/graph.cc b/src/contrib/msc/core/ir/graph.cc
index c563ccb410..2cdb326e77 100644
--- a/src/contrib/msc/core/ir/graph.cc
+++ b/src/contrib/msc/core/ir/graph.cc
@@ -1263,6 +1263,18 @@ TVM_REGISTER_GLOBAL("msc.core.MSCTensor")
return MSCTensor(name, dtype, layout, shape, alias);
});
+TVM_REGISTER_GLOBAL("msc.core.MSCTensorToJson")
+ .set_body_typed([](const MSCTensor& tensor) -> String {
+ const auto& tensor_json = tensor->ToJson();
+ std::ostringstream os;
+ dmlc::JSONWriter writer(&os);
+ tensor_json.Save(&writer);
+ return os.str();
+ });
+
+TVM_REGISTER_GLOBAL("msc.core.MSCTensorFromJson")
+ .set_body_typed([](const String& tensor_json) -> MSCTensor { return
MSCTensor(tensor_json); });
+
TVM_REGISTER_GLOBAL("msc.core.MSCJoint")
.set_body_typed([](Integer index, const String& name, const String&
shared_ref,
const String& optype, const Map<String, String>& attrs,
@@ -1293,6 +1305,11 @@ TVM_REGISTER_GLOBAL("msc.core.WeightJoint")
b_friends);
});
+TVM_REGISTER_GLOBAL("msc.core.WeightJointSetAttr")
+ .set_body_typed([](const WeightJoint& node, const String& key, const
String& value) {
+ node->attrs.Set(key, value);
+ });
+
TVM_REGISTER_GLOBAL("msc.core.MSCGraph")
.set_body_typed([](const String& name, const Array<MSCJoint>& nodes,
const Array<String>& input_names,
diff --git a/src/contrib/msc/core/ir/graph.h b/src/contrib/msc/core/ir/graph.h
index 67855deb97..fbcdeb4d0c 100644
--- a/src/contrib/msc/core/ir/graph.h
+++ b/src/contrib/msc/core/ir/graph.h
@@ -385,7 +385,7 @@ class BaseJointNode : public Object {
/*! \brief The shared_ref of node, can be changed. */
String shared_ref;
/*! \brief The attributes of node. */
- Map<String, String> attrs;
+ mutable Map<String, String> attrs;
/*! \brief The parents of node. */
Array<ObjectRef> parents;
/*! \brief The children of node. */
diff --git a/src/contrib/msc/framework/tensorrt/codegen.cc
b/src/contrib/msc/framework/tensorrt/codegen.cc
index 6efc9b26af..4cba1bdce3 100644
--- a/src/contrib/msc/framework/tensorrt/codegen.cc
+++ b/src/contrib/msc/framework/tensorrt/codegen.cc
@@ -94,7 +94,6 @@ void TensorRTCodeGen::CodeGenClassDefine() {
.func_call("malloc", "cpu_buffers[" + idx_var + "]")
.call_arg(GetTensorBytes(tensor));
};
-
stack_.line("#include \"" + graph()->name + ".h\"").line();
StartNamespace();
// start define build method
@@ -105,6 +104,12 @@ void TensorRTCodeGen::CodeGenClassDefine() {
stack_.func_arg("config", "TRTPtr<IBuilderConfig>&");
}
stack_.func_arg("logger", "TRTLogger&").func_start();
+ // save codegen before build
+ if (config()->use_tools) {
+ const auto* pf = runtime::Registry::Get("msc_tool.codegen_step");
+ ICHECK(pf != nullptr) << "Cannot find msc_tool.codegen_step func.";
+ before_build_codes_ = (*pf)(GetStepCtx(), "before_build", graph()->name,
config()->tools_tag);
+ }
if (graph()->weight_holders.size() > 0) {
stack_.assign("mWeights", "TRTUtils::LoadWeights(\"" + graph()->name +
".wts\")");
}
@@ -184,6 +189,12 @@ void TensorRTCodeGen::CodeGenClassDefine() {
.call_arg(DocUtils::ToStrDoc("use int8 to build the engine"))
.cond_end();
}
+ // save codegen after build
+ if (config()->use_tools) {
+ const auto* pf = runtime::Registry::Get("msc_tool.codegen_step");
+ ICHECK(pf != nullptr) << "Cannot find msc_tool.codegen_step func.";
+ after_build_codes_ = (*pf)(GetStepCtx(), "after_build", graph()->name,
config()->tools_tag);
+ }
// end define build method
stack_.func_end("true");
// start define test function
@@ -339,15 +350,9 @@ void TensorRTCodeGen::CodeGenMain() {
.func_call("createBuilderConfig", NullOpt, DocUtils::ToPtrDoc("builder"))
.pop_nest();
ReturnOnFail("config", "Failed to create config");
- // codegen before build
- if (config()->use_tools) {
- const auto* pf = runtime::Registry::Get("msc_tool.codegen_step");
- ICHECK(pf != nullptr) << "Cannot find msc_tool.codegen_step func.";
- const Array<String>& lines =
- (*pf)(GetStepCtx(), "before_build", graph()->name,
config()->tools_tag);
- for (const auto& l : lines) {
- stack_.line(l);
- }
+ // add codegen before build
+ for (const auto& l : before_build_codes_) {
+ stack_.line(l);
}
// build model
stack_.comment("Build model")
@@ -360,15 +365,9 @@ void TensorRTCodeGen::CodeGenMain() {
}
stack_.call_arg("logger");
ReturnOnFail("pass", "Failed to build model");
- // codegen after build
- if (config()->use_tools) {
- const auto* pf = runtime::Registry::Get("msc_tool.codegen_step");
- ICHECK(pf != nullptr) << "Cannot find msc_tool.codegen_step func.";
- const Array<String>& lines =
- (*pf)(GetStepCtx(), "after_build", graph()->name, config()->tools_tag);
- for (const auto& l : lines) {
- stack_.line(l);
- }
+ // add codegen after build
+ for (const auto& l : after_build_codes_) {
+ stack_.line(l);
}
// Set profile flag
stack_.comment("Set profile flag")
diff --git a/src/contrib/msc/framework/tensorrt/codegen.h
b/src/contrib/msc/framework/tensorrt/codegen.h
index 21b556d1ce..ea06a17f7c 100644
--- a/src/contrib/msc/framework/tensorrt/codegen.h
+++ b/src/contrib/msc/framework/tensorrt/codegen.h
@@ -84,6 +84,10 @@ class TensorRTCodeGen : public
CppCodeGen<TensorRTCodeGenConfig, TensorRTCodeGen
template <typename T>
const String ToDims(const std::vector<T>& dims, bool use_ndim = true);
const String ToDims(const Array<Integer>& dims, bool use_ndim = true);
+
+ private:
+ Array<String> before_build_codes_;
+ Array<String> after_build_codes_;
};
} // namespace msc
diff --git a/tests/python/contrib/test_msc/test_tools.py
b/tests/python/contrib/test_msc/test_tools.py
index 037507cf69..f396b81ea4 100644
--- a/tests/python/contrib/test_msc/test_tools.py
+++ b/tests/python/contrib/test_msc/test_tools.py
@@ -69,12 +69,30 @@ def _get_config(
def get_tool_config(tool_type):
+ """Get config for the tool"""
config = {}
if tool_type == ToolType.PRUNER:
config = {
"plan_file": "msc_pruner.json",
"strategys": [{"method": "per_channel", "density": 0.8}],
}
+ elif tool_type == ToolType.QUANTIZER:
+ raise NotImplementedError("Quantizer is not supported")
+ elif tool_type == ToolType.TRACKER:
+ config = {
+ "plan_file": "msc_tracker.json",
+ "strategys": [
+ {
+ "method": "save_compared",
+ "compare_to": {
+ "optimize": ["baseline"],
+ "compile": ["optimize", "baseline"],
+ },
+ "op_types": ["nn.relu"],
+ "tensor_types": ["output"],
+ }
+ ],
+ }
return {tool_type: config}
@@ -132,57 +150,67 @@ def _test_from_torch(
manager.destory()
[email protected]("tool_type", [ToolType.PRUNER])
-def test_tvm_tools(tool_type):
+def get_model_info(compile_type):
+ """Get the model info"""
+ if compile_type == MSCFramework.TVM:
+ return {
+ "inputs": [
+ {"name": "input_0", "shape": [1, 3, 224, 224], "dtype":
"float32", "layout": "NCHW"}
+ ],
+ "outputs": [{"name": "output", "shape": [1, 1000], "dtype":
"float32", "layout": "NC"}],
+ "nodes": {
+ "total": 229,
+ "input": 1,
+ "nn.conv2d": 53,
+ "nn.batch_norm": 53,
+ "get_item": 53,
+ "nn.relu": 49,
+ "nn.max_pool2d": 1,
+ "add": 16,
+ "nn.adaptive_avg_pool2d": 1,
+ "reshape": 1,
+ "msc.linear_bias": 1,
+ },
+ }
+ if compile_type == MSCFramework.TENSORRT:
+ return {
+ "inputs": [
+ {"name": "input_0", "shape": [1, 3, 224, 224], "dtype":
"float32", "layout": "NCHW"}
+ ],
+ "outputs": [{"name": "output", "shape": [1, 1000], "dtype":
"float32", "layout": ""}],
+ "nodes": {"total": 2, "input": 1, "msc_tensorrt": 1},
+ }
+ raise TypeError("Unexpected compile_type " + str(compile_type))
+
+
[email protected]("tool_type", [ToolType.PRUNER, ToolType.TRACKER])
+def test_tvm_tool(tool_type):
"""Test tools for tvm"""
- model_info = {
- "inputs": [
- {"name": "input_0", "shape": [1, 3, 224, 224], "dtype": "float32",
"layout": "NCHW"}
- ],
- "outputs": [{"name": "output", "shape": [1, 1000], "dtype": "float32",
"layout": "NC"}],
- "nodes": {
- "total": 229,
- "input": 1,
- "nn.conv2d": 53,
- "nn.batch_norm": 53,
- "get_item": 53,
- "nn.relu": 49,
- "nn.max_pool2d": 1,
- "add": 16,
- "nn.adaptive_avg_pool2d": 1,
- "reshape": 1,
- "msc.linear_bias": 1,
- },
- }
tool_config = get_tool_config(tool_type)
- _test_from_torch(MSCFramework.TVM, tool_config, model_info,
is_training=True)
+ _test_from_torch(
+ MSCFramework.TVM, tool_config, get_model_info(MSCFramework.TVM),
is_training=True
+ )
@requires_tensorrt
[email protected](
- "tool_type,use_native",
- [(ToolType.PRUNER, False)],
-)
-def test_tensorrt_tools(tool_type, use_native):
[email protected]("tool_type", [ToolType.PRUNER, ToolType.TRACKER])
+def test_tensorrt_tool(tool_type):
"""Test tools for tensorrt"""
- model_info = {
- "inputs": [
- {"name": "input_0", "shape": [1, 3, 224, 224], "dtype": "float32",
"layout": "NCHW"}
- ],
- "outputs": [{"name": "output", "shape": [1, 1000], "dtype": "float32",
"layout": ""}],
- "nodes": {"total": 2, "input": 1, "msc_tensorrt": 1},
- }
tool_config = get_tool_config(tool_type)
- if tool_type == ToolType.QUANTIZER and use_native:
+ if tool_type == ToolType.QUANTIZER:
tool_config[ToolType.QUANTIZER]["strategys"] = []
- optimize_type = MSCFramework.TENSORRT if use_native else None
+ optimize_type = MSCFramework.TENSORRT
+ else:
+ optimize_type = None
_test_from_torch(
MSCFramework.TENSORRT,
tool_config,
- model_info,
+ get_model_info(MSCFramework.TENSORRT),
is_training=False,
+ atol=5e-2,
+ rtol=5e-2,
optimize_type=optimize_type,
)
