MSC 工具测试

目录

MSC 工具测试#

%cd ..
import set_env

/media/pc/data/lxw/ai/tvm-book/doc/tutorials/msc

from tvm.contrib.msc.core.tools import ToolType

ToolType.PRUNER

'pruner'

MSCTools 与 MSCGraph 协同工作，它们决定压缩策略并控制压缩过程。MSCTools 由 RuntimeManager 管理。

from tvm.contrib.msc.core.transform import msc_transform
from tvm.contrib.msc.core.runtime import create_runtime_manager
from tvm.contrib.msc.core.tools import create_tool, MSC_TOOL

# build runtime manager from module and mscgraphs
optimized_mod, msc_graph, msc_config = msc_transform(mod, params)
rt_manager = create_runtime_manager(optimized_mod, params, msc_config)

# pruner is used for prune the model
rt_manager.create_tool(MSC_TOOL.PRUNE, prune_config)

# quantizer is used to do the calibration and quantize the model
rt_manager.create_tool(MSC_TOOL.QUANTIZE, quantize_config)

# collecter is used to collect the datas of each computational node
rt_manager.create_tool(MSC_TOOL.COLLECT, collect_config)

# distiller is used to do the knowledge distilliation
rt_manager.create_tool(MSC_TOOL.DISTILL, distill_config)

MSCGym#

MSCGym 是 MSC 中自动压缩的平台。它的作用类似于 AutoTVM，但其架构更像 OpenAI-Gym。MSCGym 从压缩过程中提取任务，然后利用代理和环境之间的交互来为每个任务找到最佳行动。要使用 MSCGym 进行自动压缩，请为工具设置 gym 配置：

{
      ...
      "quantize": {
        "strategy_file": "msc_quantize.json",
        "target": "tensorrt",
        “gym”:[
          {
            “record”:”searched_config.json”,
            “env”:{
              “strategy”:”distill_loss”
            },
            “agent”:{
              “type”:”grid_search”,
            }
          },
        ]
      },
      ...
}
目录