<jupyter_start><jupyter_text>在 Relay 中使用外部库**原作者**: [<NAME>](https://github.com/masahi), [<NAME>](https://github.com/SiNZeRo)这是简短的教程，介绍关于如何使用在 Relay 中使用外部库，如 cuDNN，或 cuBLAS。Relay 在内部使用 TVM 来生成目标特定的代码。例如，使用 cuda 后端，TVM 为用户提供的网络中的所有层生成 cuda kernel。但有时，将不同供应商开发的外部库合并到 Relay 中也是有帮助的。幸运的是，TVM 有一种透明地调用这些库的机制。对于 Relay 用户，需要做的只是适当地设置目标字符串。使用来自 Relay 的外部库之前， TVM 需要构建您想要使用的库。例如，要使用 cuDNN，在 `cmake/config.cmake` 中启用 `USE_CUDNN` 选项，必要时需要指定 cuDNN include 和库目录。首先，导入 Relay 和 TVM。<jupyter_code>import tvm
from tvm import te
import numpy as np
from tvm.contrib import graph_executor as runtime
from tvm import relay
from tvm.relay import testing
import tvm.testing<jupyter_output><empty_output><jupyter_text>创建简单网络创建非常简单的网络进行演示。它由卷积、batch normalization 和 ReLU 激活组成。<jupyter_code>out_channels = 16
batch_size = 1

data = relay.var("data", relay.TensorType((batch_size, 3, 224, 224), "float32"))
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(
    data=data, weight=weight, kernel_size=(3, 3), channels=out_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.Function(relay.analysis.free_vars(simple_net), simple_net)

data_shape = (batch_size, 3, 224, 224)
net, params = testing.create_workload(simple_net)<jupyter_output><empty_output><jupyter_text>使用 cuda 后端构建和运行像往常一样，用 cuda 后端构建和运行这个网络。通过将日志级别设置为 `DEBUG`，将 Relay graph 编译的结果转储为伪代码。<jupyter_code>import logging

logging.basicConfig(level=logging.DEBUG)  # to dump TVM IR after fusion

target = "cuda"
lib = relay.build_module.build(net, target, params=params)

dev = tvm.device(target, 0)
data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
module = runtime.GraphModule(lib["default"](dev))
module.set_input("data", data)
module.run()
out_shape = (batch_size, out_channels, 224, 224)
out = module.get_output(0, tvm.nd.empty(out_shape))
out_cuda = out.numpy()<jupyter_output>DEBUG:autotvm:Finish loading 825 records
INFO:te_compiler:Using injective.cpu for add based on highest priority (10)
/media/workspace/anaconda3/envs/mxnetx/lib/python3.10/site-packages/tvm/driver/build_module.py:263: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
  warnings.warn(
INFO:te_compiler:Using injective.cpu for sqrt based on highest priority (10)
INFO:te_compiler:Using injective.cpu for divide based on highest priority (10)
INFO:te_compiler:Using injective.cpu for multiply based on highest priority (10)
INFO:te_compiler:Using injective.cpu for expand_dims based on highest priority (10)
INFO:te_compiler:Using injective.cpu for negative based on highest priority (10)
INFO:te_compiler:Using injective.cpu for multiply based on highest priority (10)
INFO:te_compiler:Using injective.cpu for add based on highest priority (10)
INFO:te_compiler:Using injective.cpu for expand_dims based on highest priorit[...]<jupyter_text>生成的伪代码应该如下所示。```{tip}注意 bias add、batch normalization 和 ReLU 激活是如何融合到卷积核中的。```TVM 从这个表示生成单一的融合 kernel。<jupyter_code>print(lib.ir_mod["main"])
lib.function_metadata<jupyter_output><empty_output><jupyter_text>为卷积层使用 cuDNN可以用 cuDNN 来代替 cuDNN 的卷积核。为此，需要做的就是将选项 `" -libs=cudnn"` 附加到目标字符串中。<jupyter_code>net, params = testing.create_workload(simple_net)
target = "cuda -libs=cudnn"  # use cudnn for convolution
lib = relay.build_module.build(net, target, params=params)

dev = tvm.device(target, 0)
data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
module = runtime.GraphModule(lib["default"](dev))
module.set_input("data", data)
module.run()
out_shape = (batch_size, out_channels, 224, 224)
out = module.get_output(0, tvm.nd.empty(out_shape))
out_cudnn = out.numpy()<jupyter_output>DEBUG:autotvm:Finish loading 825 records
INFO:te_compiler:Using injective.cpu for add based on highest priority (10)
/media/workspace/anaconda3/envs/mxnetx/lib/python3.10/site-packages/tvm/driver/build_module.py:263: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
  warnings.warn(
INFO:te_compiler:Using injective.cpu for sqrt based on highest priority (10)
INFO:te_compiler:Using injective.cpu for divide based on highest priority (10)
INFO:te_compiler:Using injective.cpu for multiply based on highest priority (10)
INFO:te_compiler:Using injective.cpu for expand_dims based on highest priority (10)
INFO:te_compiler:Using injective.cpu for negative based on highest priority (10)
INFO:te_compiler:Using injective.cpu for multiply based on highest priority (10)
INFO:te_compiler:Using injective.cpu for add based on highest priority (10)
INFO:te_compiler:Using injective.cpu for expand_dims based on highest priorit[...]<jupyter_text>```{note}如果你使用 cuDNN, Relay 不能融合后面的层的卷积。这是因为层融合发生在 TVM 内部表示 (IR) 级别。Relay 将外部库视为黑盒，因此没有办法将它们与 TVM IR 融合。```下面的伪代码显示，cuDNN 卷积 + bias add + batch norm + ReLU 分为两个计算阶段，一个用于 cuDNN 调用，另一个用于其余的运算。<jupyter_code>lib.ir_mod<jupyter_output><empty_output><jupyter_text>验证结果可以检查两次运行的结果是否匹配。<jupyter_code>tvm.testing.assert_allclose(out_cuda, out_cudnn, rtol=1e-5)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11195016046818954, 'fraction_numerical': 0.019444968850292618, 'mean_word_length': 7.614634146341463, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 29, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '421478', 'license': 'Apache-2.0', 'n_tokens_mistral': 2257, 'n_tokens_neox': 2248, 'n_words': 429, 'path': 'xinetzone/docs/how_to/work_with_relay/using_external_lib.ipynb', 'repo_name': 'xinetzone/tvm', 'size': '5311'}