Skip to content

Commit

Permalink
remove gym test
Browse files Browse the repository at this point in the history
  • Loading branch information
@Archermmt
Archermmt committed Jan 2, 2024
1 parent 1fb2793 commit 0c9bef7
Show file tree
Hide file tree
Showing 3 changed files with 8 additions and 76 deletions.
5 changes: 1 addition & 4 deletions tests/python/contrib/test_msc/test_manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
""" Test Managers in MSC. """

import json
import datetime
import pytest
import torch

Expand All @@ -36,9 +35,7 @@
def _get_config(model_type, compile_type, inputs, outputs, atol=1e-1, rtol=1e-1):
"""Get msc config"""

path = "test_manager_{}_{}_{}".format(
model_type, compile_type, datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
)
path = "test_manager_{}_{}".format(model_type, compile_type)
return {
"workspace": msc_utils.msc_dir(path),
"verbose": "critical",
Expand Down
13 changes: 5 additions & 8 deletions tests/python/contrib/test_msc/test_runner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@

""" Test Runners in MSC. """

import datetime
import pytest
import numpy as np

Expand Down Expand Up @@ -84,12 +83,10 @@ def _test_from_torch(runner_cls, device, is_training=False, atol=1e-1, rtol=1e-1

torch_model = _get_torch_model("resnet50", is_training)
if torch_model:
path = "test_runner_torch_{}_{}_{}".format(
runner_cls.__name__, device, datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
)
path = "test_runner_torch_{}_{}".format(runner_cls.__name__, device)
workspace = msc_utils.set_workspace(msc_utils.msc_dir(path))
log_path = workspace.relpath("MSC_LOG", keep_history=False)
msc_utils.set_global_logger("info", log_path)
msc_utils.set_global_logger("critical", log_path)
input_info = [([1, 3, 224, 224], "float32")]
datas = [np.random.rand(*i[0]).astype(i[1]) for i in input_info]
torch_datas = [torch.from_numpy(d) for d in datas]
Expand All @@ -114,7 +111,7 @@ def test_tvm_runner_cpu():

@tvm.testing.requires_gpu
def test_tvm_runner_gpu():
"""Test runner for tvm on gpu"""
"""Test runner for tvm on cuda"""

_test_from_torch(TVMRunner, "cuda", is_training=True)

Expand Down Expand Up @@ -144,10 +141,10 @@ def test_tensorflow_runner():

tf_graph, graph_def = _get_tf_graph()
if tf_graph and graph_def:
path = "test_runner_tf_" + str(datetime.datetime.now().strftime("%Y%m%d_%H%M%S"))
path = "test_runner_tf"
workspace = msc_utils.set_workspace(msc_utils.msc_dir(path))
log_path = workspace.relpath("MSC_LOG", keep_history=False)
msc_utils.set_global_logger("info", log_path)
msc_utils.set_global_logger("critical", log_path)
data = np.random.uniform(size=(1, 224, 224, 3)).astype("float32")
out_name = "MobilenetV2/Predictions/Reshape_1:0"
# get golden
Expand Down
66 changes: 2 additions & 64 deletions tests/python/contrib/test_msc/test_tools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
""" Test Tools in MSC. """

import json
import datetime
import pytest
import torch

Expand Down Expand Up @@ -46,12 +45,7 @@ def _get_config(
):
"""Get msc config"""

path = "test_tool_{}_{}".format(model_type, compile_type)
for t_type, config in tools_config.items():
path = path + "_" + str(t_type)
if "gym_configs" in config:
path = path + "_gym"
path = path + "_" + str(datetime.datetime.now().strftime("%Y%m%d_%H%M%S"))
path = "_".join(["test_tools", model_type, compile_type] + list(tools_config.keys()))
return {
"workspace": msc_utils.msc_dir(path),
"verbose": "critical",
Expand All @@ -76,7 +70,7 @@ def _get_config(
}


def get_tool_config(tool_type, use_distill=False, use_gym=False):
def get_tool_config(tool_type, use_distill=False):
"""Get config for the tool"""

config = {}
Expand All @@ -85,23 +79,6 @@ def get_tool_config(tool_type, use_distill=False, use_gym=False):
"plan_file": "msc_pruner.json",
"strategys": [{"method": "per_channel", "density": 0.8}],
}
if use_gym:
config["gym_configs"] = [
{
"env": {
"executors": {
"action_space": {
"method": "action_prune_density",
"start": 0.4,
"end": 0.8,
"step": 0.4,
}
},
"max_tasks": 3,
},
"agent": {"agent_type": "search.grid", "executors": {}},
}
]
elif tool_type == ToolType.QUANTIZER:
# pylint: disable=import-outside-toplevel
from tvm.contrib.msc.core.tools.quantize import QuantizeStage
Expand Down Expand Up @@ -139,23 +116,6 @@ def get_tool_config(tool_type, use_distill=False, use_gym=False):
},
],
}
if use_gym:
config["gym_configs"] = [
{
"env": {
"executors": {
"action_space": {
"method": "action_quantize_scale",
"start": 0.8,
"end": 1.2,
"step": 0.2,
}
},
"max_tasks": 3,
},
"agent": {"agent_type": "search.grid", "executors": {}},
}
]
elif tool_type == ToolType.TRACKER:
config = {
"plan_file": "msc_tracker.json",
Expand Down Expand Up @@ -302,17 +262,6 @@ def test_tvm_distill(tool_type):
)


@tvm.testing.requires_gpu
@pytest.mark.parametrize("tool_type", [ToolType.PRUNER, ToolType.QUANTIZER])
def test_tvm_gym(tool_type):
"""Test tools for tvm with gym"""

tool_config = get_tool_config(tool_type, use_gym=True)
_test_from_torch(
MSCFramework.TVM, tool_config, get_model_info(MSCFramework.TVM), is_training=True
)


@requires_tensorrt
@pytest.mark.parametrize(
"tool_type",
Expand Down Expand Up @@ -349,16 +298,5 @@ def test_tensorrt_distill(tool_type):
)


@requires_tensorrt
@pytest.mark.parametrize("tool_type", [ToolType.PRUNER])
def test_tensorrt_gym(tool_type):
"""Test tools for tensorrt with gym"""

tool_config = get_tool_config(tool_type, use_gym=True)
_test_from_torch(
MSCFramework.TENSORRT, tool_config, get_model_info(MSCFramework.TENSORRT), is_training=False
)


if __name__ == "__main__":
tvm.testing.main()

0 comments on commit 0c9bef7

Please sign in to comment.