重庆专业建网站,wordpress修改 id,台州网站制作台州网站建设,全英文外贸网站建设Apache TVM 是一个端到端的深度学习编译框架#xff0c;适用于 CPU、GPU 和各种机器学习加速芯片。更多 TVM 中文文档可访问 → Apache TVM 中文站tvm.hyper.ai/
作者#xff1a;Hua Jiang
本教程介绍如何将「Pipeline Executor」与 Relay 配合使用。
import tvm
from t…Apache TVM 是一个端到端的深度学习编译框架适用于 CPU、GPU 和各种机器学习加速芯片。更多 TVM 中文文档可访问 → Apache TVM 中文站tvm.hyper.ai/
作者Hua Jiang
本教程介绍如何将「Pipeline Executor」与 Relay 配合使用。
import tvm
from tvm import te
import numpy as np
from tvm.contrib import graph_executor as runtime
from tvm.relay.op.contrib.cutlass import partition_for_cutlass
from tvm import relay
from tvm.relay import testing
import tvm.testing
from tvm.contrib.cutlass import finalize_modulesimg_size 8创建一个简单的网络这个网络也可以是一个预训练的模型。
创建一个由 convolution、batch normalization、dense 和 ReLU activation 组成的网络用于演示。
def get_network():out_channels 16batch_size 1data relay.var(data, relay.TensorType((batch_size, 3, img_size, img_size), float16))dense_weight relay.var(dweight, relay.TensorType((batch_size, 16 * img_size * img_size), float16))weight relay.var(weight)bn_gamma relay.var(bn_gamma)bn_beta relay.var(bn_beta)bn_mmean relay.var(bn_mean)bn_mvar relay.var(bn_var)simple_net relay.nn.conv2d(datadata, weightweight, kernel_size(3, 3), channelsout_channels, padding(1, 1))simple_net relay.nn.batch_norm(simple_net, bn_gamma, bn_beta, bn_mmean, bn_mvar)[0]simple_net relay.nn.relu(simple_net)simple_net relay.nn.batch_flatten(simple_net)simple_net relay.nn.dense(simple_net, dense_weight)simple_net relay.Function(relay.analysis.free_vars(simple_net), simple_net)data_shape (batch_size, 3, img_size, img_size)net, params testing.create_workload(simple_net)return net, params, data_shapenet, params, data_shape get_network()将网络拆分成两个子图。
这个来自单元测试的名为「graph_split」的函数只是一个例子。用户可以创建自定义逻辑来拆分计算图。
import inspect
import ostutorial_dir os.path.dirname(inspect.getfile(lambda: None))
os.sys.path.append(os.path.join(tutorial_dir, ../../../tests/python/relay))
from test_pipeline_executor import graph_split将网络拆分成两个子图。
split_config [{op_name: nn.relu, op_index: 0}]
subgraphs graph_split(net[main], split_config, params)生成的子图如下所示。 #subgraphs[0])def main(%data: Tensor[(1, 3, img_size, img_size), float16]) {%0 nn.conv2d(%data, meta[relay.Constant][0] /* tyTensor[(16, 3, 3, 3), float16] */, padding[1, 1, 1, 1], channels16, kernel_size[3, 3]) /* tyTensor[(1, 16, img_size, img_size), float16] */;%1 nn.batch_norm(%0, meta[relay.Constant][1] /* tyTensor[(16), float16] */, meta[relay.Constant][2] /* tyTensor[(16), float16]*/, meta[relay.Constant][3] /* tyTensor[(16), float16] */, meta[relay.Constant][4] /* tyTensor[(16), float16] */) /* ty(Tensor[(1,16, img_size, img_size), float16], Tensor[(16), float16], Tensor[(16), float16]) */;%2 %1.0;nn.relu(%2) /* tyTensor[(1, 16, img_size, img_size), float16] */}#subgraphs[1]def main(%data_n_0: Tensor[(1, 16, 8, 8), float16] /* tyTensor[(1, 16, 8, 8), float16] */) {%0 nn.batch_flatten(%data_n_0) /* tyTensor[(1, 1024), float16] */;nn.dense(%0, meta[relay.Constant][0] /* tyTensor[(1, 1024), float16] */, unitsNone) /* tyTensor[(1, 1), float16] */}用 cutlass target 构建子图。
cutlass tvm.target.Target({kind: cutlass,sm: int(tvm.target.Target(cuda).arch.split(_)[1]),use_3xtf32: True,split_k_slices: [1],profile_all_alignments: False,find_first_valid: True,use_multiprocessing: True,use_fast_math: False,tmp_dir: ./tmp,},hosttvm.target.Target(llvm),
)def cutlass_build(mod, target, paramsNone, target_hostNone, mod_namedefault):target [target, cutlass]lib relay.build_module.build(mod, targettarget, paramsparams, target_hosttarget_host, mod_namemod_name)return lib使用 pipeline executor 在 pipeline 中运行两个子图。
在 cmake 中将 USE_PIPELINE_EXECUTOR 和 USE_CUTLASS 设置为 ON。
from tvm.contrib import graph_executor, pipeline_executor, pipeline_executor_build创建子图 pipeline 配置。将子图模块与 target 关联起来。使用 CUTLASS BYOC 构建第二个子图模块。
mod0, mod1 subgraphs[0], subgraphs[1]
# 将 cutlass 作为 codegen。
mod1 partition_for_cutlass(mod1)获取 pipeline executor 配置对象。
pipe_config pipeline_executor_build.PipelineConfig()设置子图模块的编译 target。
pipe_config[mod0].target llvm
pipe_config[mod0].dev tvm.cpu(0)将第二个子图模块的编译 target 设置为 cuda。
pipe_config[mod1].target cuda
pipe_config[mod1].dev tvm.device(cuda, 0)
pipe_config[mod1].build_func cutlass_build
pipe_config[mod1].export_cc nvcc
# 通过连接子图模块创建 pipeline。
# 全局输入将被转发到第一个名为 mod0 的模块的输入接口
pipe_config[input][data].connect(pipe_config[mod0][input][data])
# mod0 的第一个输出会转发到 mod1 的输入接口
pipe_config[mod0][output][0].connect(pipe_config[mod1][input][data_n_0])
# mod1 的第一个输出将是第一个全局输出。
pipe_config[mod1][output][0].connect(pipe_config[output][0])pipeline 配置如下 print(pipe_config)Inputs|data: mod0:dataoutput|output(0) : mod1.output(0)connections|mod0.output(0)- mod1.data_n_0构建 pipeline executor。
with tvm.transform.PassContext(opt_level3):pipeline_mod_factory pipeline_executor_build.build(pipe_config)输出结果
/workspace/python/tvm/driver/build_module.py:267: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, hosttarget_host) instead.target_host parameter is going to be deprecated. 将参数配置导出到一个文件中。
directory_path tvm.contrib.utils.tempdir().temp_dir
os.makedirs(directory_path, exist_okTrue)
config_file_name pipeline_mod_factory.export_library(directory_path)使用 load 函数创建和初始化 PipelineModule。
pipeline_module pipeline_executor.PipelineModule.load_library(config_file_name)运行 pipeline executor。
分配输入数据。
data np.random.uniform(-1, 1, sizedata_shape).astype(float16)
pipeline_module.set_input(data, tvm.nd.array(data))以 pipeline 模式运行两个子图异步或同步获取输出。以下示例为同步获取输出。
pipeline_module.run()
outputs pipeline_module.get_output()使用 graph_executor 进行验证。
用 graph_executor 依次运行这两个子图得到输出。
target llvm
dev0 tvm.device(target, 0)
lib0 relay.build_module.build(mod0, target, paramsparams)
module0 runtime.GraphModule(lib0[default](dev0))
cuda tvm.target.Target(cuda, hosttvm.target.Target(llvm))
lib1 relay.build_module.build(mod1, [cuda, cutlass], paramsparams)
lib1 finalize_modules(lib1, compile.so, ./tmp)dev1 tvm.device(cuda, 0)module1 runtime.GraphModule(lib1[default](dev1))module0.set_input(data, data)
module0.run()
out_shape (1, 16, img_size, img_size)
out module0.get_output(0, tvm.nd.empty(out_shape, float16))
module1.set_input(data_n_0, out)
module1.run()
out_shape (1, 1)
out module1.get_output(0, tvm.nd.empty(out_shape, float16))输出结果
/workspace/python/tvm/driver/build_module.py:267: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, hosttarget_host) instead.target_host parameter is going to be deprecated. 验证结果。
tvm.testing.assert_allclose(outputs[0].numpy(), out.numpy())下载 Python 源代码using_pipeline_executor.py
下载 Jupyter Notebookusing_pipeline_executor.ipynb
