PyTorch 与 TVM 的桥梁
（更偏重代码的版本交叉发布在更 PyTorch 相关的 Lernapparat 上，Jupyter Notebook 教程可以在 github 上找到。）
人工智能最令人着迷的应用之一是在自然语言处理领域。像 BERT 或 GPT-2 及其变体这样的模型似乎能够掌握足够的文本信息，从而以一种需要仔细辨认才能识别为胡言乱语的方式继续下去。
这些模型属于一类称为Transformer的神经网络架构。实现它们最受欢迎的库之一是 HuggingFace transformers 库。
但是，与卷积模型或 LSTM 相比，我们对卷积模型或 LSTM 进行了大量优化实现，而 Transformer 模型的情况并非如此。因此，在这里我们探索 TVM 如何填补这一空白。我们将分两步进行：
- 首先，我们研究 BERT 推理，并在 TVM 上对其进行调优。
- 其次，我们开始更深入地探索如何在 PyTorch 中使用 TVM 进行训练。鉴于其实验性质，我们更侧重于可行性，而不是这一部分的性能。
使用 TVM 优化 BERT 推理
那么，我们如何将 BERT 从 transformer 库转移到 TVM 呢？
值得庆幸的是，transformers 支持使用 PyTorch JIT 追踪其模型。我们使用他们的 相关教程，特别是直到我们获得追踪模型的部分。
在我的 AMD Radeon VII 上，PyTorch 追踪模型对示例输入进行 100 次运行大约需要 0.65-0.7 秒，这意味着每次运行 6.5-7 毫秒。我们可以尝试看看是否可以使用 TVM 获得更快的速度。将我们的模型转换为 TVM 非常轻松：
shape_list = [(i.debugName().split('.')[0], i.type().sizes()) for i in list(traced_model.graph.inputs())[1:]]
mod_bert, params_bert = tvm.relay.frontend.pytorch.from_pytorch(traced_model,
shape_list, default_dtype="float32")
会有一些关于找不到 dtype 信息的警告，但总体进展顺利！我们现在可以构建并运行它。构建过程遵循标准的 TVM 流程。我们还将 PyTorch（cpu）张量转换为 TVM 数组。
target = 'rocm -model=gfx906' # use what matches your GPU
target_host = 'llvm'
ctx = tvm.context(target)
tt_a = tvm.nd.array(tokens_tensor.numpy(), ctx)
st_a = tvm.nd.array(segments_tensors.numpy(), ctx)
tvm.relay.backend.compile_engine.get().clear() # just to be sure, see https://github.com/apache/incubator-tvm/pull/5724
with tvm.transform.PassContext(opt_level=3):
graph, lib, params = tvm.relay.build(mod_bert,
target=target,
target_host=target_host,
params=params_bert)
module = tvm.contrib.graph_runtime.create(graph, lib, ctx)
这将多次警告我们：
WARNING:autotvm:Cannot find config for ... batch_matmul.cuda .... A fallback configuration is used, which may bring great performance regression.
哎呀，可能会带来巨大的性能下降。让我们看看。
但首先我们运行模型，看看输出是否匹配：
(8.583069e-06, 8.493662e-07)
看起来不错。请记住，我们正在 float32 中进行计算，因此 $10^{-6}$ 左右是一个不错的结果。
构建我们的模型并设置参数后，我们像这样计时我们的模型：
def x():
for i in range(100):
module.run()
ctx.sync()
x()
%timeit x()
哎哟，每次 100 次运行需要 6.65 秒，即每次运行模型 67 毫秒。这确实很慢。但是警告说，这是因为它找不到（调优的）配置。那么让我们调优任务。
调优确实需要半天左右的时间（我基本上是在遵循 TVM 调优教程，使用 autotvm 进行 ResNet 调优。）
完成此操作后，我们可以再次构建模型，这次使用新配置。这次我们应该看不到任何关于缺少配置的注释。现在它在每次运行 6.5-7 毫秒的范围内，与 PyTorch 相似。这是我们从对算子的非常基本的优化中获得的结果。但是，我们可以稍微进一步推进。
为了了解如何操作，让我们深入了解 BERT 建模和 TVM。
如果您不想了解全部细节，请跳过下一节，向下滚动到结果。我应该补充一点，我希望本教程的调优部分在不久的将来会过时，因为您会获得更好的开箱即用速度，或者至少在一些初始调优之后。因此，如果您在这里和结果之间没有看到加速，那是因为我在提交补丁方面做了功课。
BERT 模型
让我们仔细看看 BERT 中发生了什么。
像许多深度学习模型一样，BERT 带有一些序言（词汇嵌入）和尾声（池化），主体部分被组织成看起来相似的块，这里我们有 12 个 BertLayer 模块。attention_mask 只是为了防止 BERT 在处理问题时查看答案。
因此，让我们放大并详细查看 BertLayer，因为这最终是我们需要快速完成的事情。正如我们在网络图中看到的那样，BertLayer 模块的主要部分是子模块 BertSelfAttention。
现在，BertSelfAttention 捕获了著名的自注意力机制，这是 Transformer 模型的标志。（我强烈推荐 Sascha Rush 的 Annotated Transformer，它是一个详细的演练。）
在显微镜下观察 BertLayer
如果我们想深入了解细节，我们应该希望单独运行 BertLayer。我们获取 BertLayer 的输入（请参阅 Notebook 了解如何操作），并将单个 BertLayer 转换为 TVM，就像我们对整个模型所做的那样。
为了查看 TVM 模块，我们定义了一个小的可视化助手（大致基于 TVM PR#4370）。
import graphviz
def visualize(expr, collapse_small=True, node_attr_dict = {}):
def collect_ops(node):
ops = set()
def visitor(e):
if isinstance(e, tvm.ir.Op):
ops.add(e.name)
tvm.relay.analysis.post_order_visit(node, visitor)
return ops
# node_dict maps a Relay node to an index (node ID)
def _traverse_expr(node, node_dict):
if node in node_dict:
return
node_dict[node] = len(node_dict)
node_dict = {}
tvm.relay.analysis.post_order_visit(expr, lambda x: _traverse_expr(x, node_dict))
relayviz_nodes = []
dot = graphviz.Digraph(format='svg', )
dot.attr('node', shape = 'box')
def to_str(node):
if isinstance(node, tvm.relay.Constant):
return repr(node).lstrip('Constant(')[:-1]
else:
raise NotImplementedError("to_str:" + repr(node))
def is_small_const(c):
if not (collapse_small and isinstance(c, tvm.relay.Constant)):
return False
if isinstance(c.data, tvm.runtime.ndarray.NDArray):
return numpy.prod(c.data.shape) < 10
return True
# Sort by node ID
for node, node_id in sorted(node_dict.items(), key=lambda x: x[1]):
if isinstance(node, tvm.relay.Function):
dot.node(str(node_id), 'Function', **node_attr_dict.get(node, {}))
dot.edge(str(node_dict[node.body]), str(node_id))
elif isinstance(node, tvm.relay.Var):
if node.type_annotation is not None:
if hasattr(node.type_annotation, 'shape'):
shape = tuple([int(x) for x in node.type_annotation.shape])
dtype = node.type_annotation.dtype
typstr = 'Tensor[{}, {}]'.format(shape, dtype)
else:
typstr = str(node.type_annotation)
else:
typstr = '?'
d = dict(shape = 'ellipse')
d.update(node_attr_dict.get(node, {}))
dot.node(str(node_id),
'{}: {}'.format(
node.name_hint, typstr
), **d)
elif isinstance(node, tvm.relay.Tuple):
dot.node(str(node_id), 'Tuple[...])', **node_attr_dict.get(node, {}))
for field in node.fields:
dot.edge(str(node_dict[field]), str(node_id))
elif isinstance(node, tvm.relay.Constant):
if not is_small_const(node): # small consts are shown in ops
dot.node(str(node_id), 'Constant({}, {})'.format(node.data.shape, node.data.dtype),
**node_attr_dict.get(node, {}))
elif isinstance(node, tvm.relay.Call):
args_with_edge = []
arg_str_list = []
for arg in node.args:
if is_small_const(arg):
arg_str_list.append(to_str(arg))
else:
arg_str_list.append('·')
args_with_edge.append(arg)
arg_str = ', '.join(arg_str_list)
if isinstance(node.op, tvm.ir.Op):
name = node.op.name
attrs = {k:getattr(node.attrs, k) for k in node.attrs.keys()} if hasattr(node.attrs, 'keys') else {}
#attrs = inspect.getmembers(node.attrs)
attr_str_list = [k+'='+(str(v) if len(str(v))<20 else "...") for k, v in attrs.items()]
if attr_str_list:
attr_str = '| '+ ', '.join(attr_str_list)
else:
attr_str = ''
else:
ops = collect_ops(node)
if ops:
name = '_'.join(ops)
else:
name = '...'
attr_str = ''
s = f'{name}({arg_str}{attr_str})'
dot.node(str(node_id), s, **node_attr_dict.get(node, {}))
for arg in args_with_edge:
dot.edge(str(node_dict[arg]), str(node_id))
elif isinstance(node, tvm.ir.Op):
# dot.node(str(node_id), 'Op {}'.format(node.name))
pass # covered in call
elif isinstance(node, tvm.relay.TupleGetItem):
dot.node(str(node_id), 'TupleGetItem(idx={})'.format(node.index), **node_attr_dict.get(node, {}))
dot.edge(str(node_dict[node.tuple_value]), str(node_id))
elif isinstance(node, tvm.relay.Let):
dot.node(str(node_id), 'Let(XX)', **node_attr_dict.get(node, {}))
dot.edge(str(node_dict[node.value]), str(node_id))
dot.edge(str(node_id), str(node_dict[node.var]))
else:
raise RuntimeError(
'Unknown node type. node_id: {}, node: {}'.format(node_id, type(node)))
return dot
让我们在我们的主函数上运行它。出于某种原因（好吧，可能是为了完全通用），PyTorch 转换器会将 Linear 层转换为 batch_matmul，而不是仅仅 dense。我们稍后会回到这一点。由于 TVM 的 batch_matmul 在两个操作数上都将收缩轴放在最后（与 PyTorch 不同），因此也有相当多的转置操作。
visualize(mod['main'])
除了我们命名的输入之外，我们还看到许多未命名（编号）的变量。这些是神经网络参数。
让我们编译我们的模型。
与完整模型一样，在检查它是否计算相同的量后，我们可以运行并计时我们的子模块。
100 次运行需要 20.2 毫秒。这里的粗略计算是，对于 PyTorch 中的 BertLayer，我们在这个层中花费大约 0.2 毫秒，因此在 12 层上花费大约 2.4 毫秒——不是主要的，但占 6-7 毫秒总运行时间的可观部分。让我们与 TVM 进行比较。（一个好的规则是永远不要在没有测量的情况下进行优化。）
同样，TVM 的时钟显示 100 次运行为 18.2 毫秒。因此，在这里我们再次与 PyTorch 大致相当。
我们从图片中看到的一件事是输入被重塑了三次。有一个 TVM 优化 pass 称为公共子表达式消除 (CSE)，它组合了三个重塑。（不久前，这没有成功，因为它具有不同的形状参数，但这自那时以来已由 TVM 开发人员在动态到静态转换 pass 中解决。）此外，模型参数也被重塑和转置。我们也可以摆脱它吗？是的。为此，我们首先需要绑定参数，即将它们放入模型中。然后参数变成了常量而不是输入节点。使用 Foldconstant pass，我们可以通过 transpose 和 reshape 传播常量，使它们更接近 matmul。
在这三个之后（当我们编译 relay 模型时，TVM 会执行这些操作），我们的模型看起来像这样：
现在有一个有趣的技巧。将具有相同输入的三个批次 matmul 合并为单个 batch_matmul 更有效。我们在 TVM PR 5791 中实现了一个执行此操作的 pass。因此，让我们调用它，并进行另一个常量折叠 pass。
new_mod = tvm.relay.transform.CombineParallelBatchMatmul()(new_mod)
new_mod = tvm.relay.transform.FoldConstant()(new_mod)
visualize(new_mod["main"])
太棒了。在检查我们是否仍然获得相同的结果之后。我们可以再次计时：100 次运行为 25.2 毫秒。它又有点慢了，因为我们需要为新形状进行调优。调优后，我们 100 次运行为 12.6 毫秒，因此我们从大约 0.2 毫秒降至大约 0.13-0.15 毫秒，这是一个不错的加速。通过我们粗略的计算，这应该从总运行时间中减少 0.6-0.8 毫秒，或者在 5%-10% 之间。让我们检查一下。
优化后整体 BERT 模型的结果
让我们定义一个函数，将上面的优化 pass 组合起来，并在整个 BERT 模型上运行它。我们执行与上面相同的操作。
我们得到 100 次运行为 624 毫秒。太棒了，我们从 PyTorch 中的 6.5-7 毫秒降至 TVM 中的约 6.2 毫秒。这是一个 5%-10% 的加速。请注意，我们只采用了一个特定的、不太大的形状。更认真的分析会考虑更多的问题形状。
我们可能会进一步推进——例如，通过处理重塑来融合批次 matmul 之后的加法，但我们现在就此为止。我们也将受益于 TVM 的进一步改进，因此随着时间的推移，基准测试的改进情况将很有趣。特别是，即将推出的 Ansor 调优机制看起来很有希望。
幕后一瞥
比较模型的实现
正如您所看到的，我始终将 PyTorch 与 TVM 输出进行比较，以查看它们是否良好。此外，当我调查一些内部层时，我抓取了该层的输入以转换为并馈送到 TVM 模型。我确实相信这是一种非常有效的技术。
但是，有时很难评估结果之间的偏差是来自数值精度还是来自某处的错误。当我最初转换模型时，SelfAttention 子模块输出被 TVM 模型复制到大约 1e-6。但是，BertLayer 转换的结果约为 1-e3。我不太清楚这可能是由于累积的数值误差还是某处的实质性偏差。（事实证明是 GELU 激活，它被转换为 FastGELU。）
在这种情况下，我喜欢做的一件事是跳转到双精度并检查那里。数值误差应该变得小得多，而其他偏差将保持相同的数量级。使用 PyTorch 前端，如果您将 default_dtype="float64" 传递给转换函数，则可以在 PyTorch 端跟踪转换为 float64 的模型。
运行模块并与 PyTorch 进行比较现在应该有 1e-14 左右的偏差。
TVM 中为方便此用例而进行的改进
在像这里显示的那样工作之前，我们必须弥合一些差距（但最近的 git 签出将包括所有这些差距）：
- TVM PyTorch 转换器不支持 fp32 以外的输入。我们 实现了改进的转换，现在也包含在 TVM upstream 中。
- 工作负载操作 batch_matmul 的 TVM 调度（即计算的组织）是固定的，并且非常慢（类似于现在在没有调优调度的情况下运行）。因此，我们 实现了一个可调优的调度。
- PyTorch 转换器生成批次 matmul 操作（它也可能被更改为生成密集层）。但是正如我们所看到的，更大的速度优势之一是组合查询键和值线性层，因此我们 实现了融合批次 matmul 操作。
- 在比较计算结果时，我们注意到 GELU 函数被转换为其 FastGELU 变体。我们修复了这个问题。（TVM 中有一个快速数学优化 pass，它对误差函数进行了一些替换，尽管我们没有检查它是否为用误差函数表示的 GELU 生成 FastGELU。）
- TVM 最初（并且在某种程度上仍然是）专注于静态形状。最近，它尝试了动态操作。动态重塑——获取目标形状的参数——是这些实验的早期阶段，但正如上面所看到的，它阻止了批次 matmul 的融合，因为公共子表达式消除 pass 没有检测到它可以合并相同的输入重塑。最近这种情况有所改善。
使用 TVM 计算训练 Pytorch 模型
在第二部分中，我们想看看是否可以在 PyTorch 中训练 BERT 时使用 TVM。当然，这开启了一个全新的难题，因为我们需要处理自动微分。当我们继续上面的主题并以 BertLayer 为例时，我们的方法通常代表了非平凡的模块。我们希望在训练期间将计算转移到 TVM。
因此，用户可以获取一个（可追踪的）模块并执行：
add_tvm_dispatch(module, sample_input)
然后，如果她使用与 sample_input 相同形状的输入调用模块，她将获得由 TVM 计算的输出（当然，作为 PyTorch 张量），如果不是，它将只使用常规 forward。
但是，所以我们已经暗示了坏消息：在这一部分中，我们将看到如何做这些事情。我们尚未实现巨大的加速。
但说够了，让我们直接深入！同样，我们通过运行来自 transformer Bert 模型的追踪 BertLayer 通过 tvm.relay.frontend.from_pytorch 获得我们的 relay 模型。
我们将在中间做的一件事是从 PyTorch 中的模块化接口（带有命名参数）转移到函数式接口（这正是 TVM 可以为我们做的事情）。我们为此要做的第一件事是安排函数参数的顺序，以便我们可以使用它——即首先是模块的直接输入，然后是参数，顺序与 PyTorch 使用它们的顺序相同。在此操作之后，我们在 TVM 中的 BertLayer 看起来像这样：
与 BERT 推理一样，我们希望运行一些优化 pass。
但我们也有一些新的转换：
- 自动微分的一个特殊性是它会使用大量的
..._like操作来广播或“取消广播”（求和是关于自动微分的广播的对偶）事物。但这意味着您现在有两个张量参数，即使后者实际上不需要梯度。ZappLike用相应的函数替换这些操作，这些函数采用形状参数而不是。 - 另一件事是导数的“生根”。TVM 生成一个张量，其中包含与我们函数的返回值形状相同的所有 1 作为链式法则的起点。然后将这些张量乘以我们操作的导数。但是与 1 相乘并没有做太多事情，所以我们将其删除。同样，TVM 将变量（输入）的梯度初始化为相同形状的零。如果未使用它，则梯度将为零，但如果使用了，则“真实梯度”将添加到该零。但是也可以消除加零。这些由 ZeroZapp 和 OneZapp 处理。
- TVM 没有
LayerNorm（或BatchNorm或其他）的训练变体。因此，我们实现了一个 pass 来详细说明计算。 - TVM 也没有训练 dropout。这里的问题有点难以解决，因为 TVM 目前没有随机数。我们改为用一个构造来替换 dropout，该构造采用随机伯努利抽样（0/1 值）并使用该抽样模拟 dropout。我们的想法是，我们将使用 PyTorch 为我们生成此掩码。这样做的好处是（如果我们以与 PyTorch 相同的顺序生成 dropout 掩码），我们将获得完全相同的结果。
如上文暗示的那样，TVM 的梯度计算假定它是计算中的最后一个元素（上面讨论的 ones-Tensors）。这与 PyTorch 的模块化视图不太匹配，后者期望为每个给定的输出提供一个 grad_out。幸运的是，这在计算上等同于乘以 grad out 和求和，因此我们用它来修改我们的函数。我们希望具有灵活性，因此我们允许返回单个张量的函数和返回张量元组的函数。
应用这些修改后，我们的模型看起来像这样：
最后，我们可以求梯度。由于我们获得了许多 let 节点，我们使用 ToGraphNormalForm pass 将其带到正常形式。TVM 的梯度计算返回一个函数，该函数具有与原始函数相同的参数（在我们的例子中，用 grad_out 和 dropout 进行了修改），然后返回原始返回值的元组和一个包含所有输入的梯度的元组。我们做的第一件事是删除我们不需要的 grad_out 和 dropout 的所有梯度。然后我们运行我们的简化 pass。
所以这是我们现在拥有的前向和后向图：
但是在 PyTorch 中，我们首先计算前向，然后计算后向，因此我们必须取出锯子并拆分我们的图。困难的问题之一是如何处理为前向和后向计算的事物。这是一个难题，与 MinCut 问题有关。
我们的极端选择可能是：
- 可以只保留输入并根据需要重新计算所有内容。
- 如果我们有一个标量输出，我们可以计算梯度并在后向时乘以后面层的导数。（损失函数可能会这样做。）但是，这不适用于非标量张量输出。
我们将执行以下操作：我们正常计算前向，但我们保留后向中将要使用的所有内容。不幸的是，这太多了，这很可能是我们看不到端到端加速的原因。我们将在下面讨论一些潜在的启发式方法。
我们在这里使用着色。首先，我们将前向计算的所有节点着色为红色。然后，我们遍历梯度计算，然后将其从后向需要的节点着色为蓝色。这让我们有机会展示我们可视化中的属性支持。
一点（PyTorch）术语：当我们有一个函数Layer : x ↦ y，后跟一些Loss: y ↦ l ∈ ℝ时，后向是BackwardOfLayer : grad_out ↦ grad_in，其中grad_out = dl/dy 和 *grad_in = dl/dx`。
为了如上所述拆分函数，我们收集蓝色节点以进行捕获 - 但常量将被复制，输入（Var 节点）需要单独处理。现在我们可以拆分后向，用变量替换所有蓝色节点。
接下来，我们获取前向并对其进行修改，使其也返回所需的中间值。然后前向看起来像这样：
TVM 无法返回嵌套元组，因此我们在函数中展平输出。我们再次区分张量值函数和元组值函数（即那些可能返回多个张量的函数）。
最后，我们可以让 TVM 发挥其魔力并编译我们的函数，例如编译为 gr_only_compiled_module 和 fw_and_cap_compiled_module。是时候试一试了。我们定义了便利函数以在 PyTorch 和 TVM 之间移动张量，并将模型参数作为 TVM 字典获取。
def tensor_to_tvm(t):
return tvm.nd.from_dlpack(torch.utils.dlpack.to_dlpack(t))
def tensor_from_tvm(a):
return(torch.utils.dlpack.from_dlpack(a.to_dlpack()))
model_params_tvm = {k: tensor_to_tvm(v) for k, v in pytorch_model.state_dict().items()}
同样，我们在 PyTorch 和 TVM 中获取 GPU 上的输入。
我们需要处理 dropout。事实证明，我们对三个 dropout 随机抽样的记录与模型中的 dropout 以相同的顺序发生。我们对计算图进行了深度优先搜索以找到它们，如果 dropout 的值在图中连接而不是在独立分支上，则这将是 PyTorch 绘制矩阵的顺序。如果不是，祝您好运调整顺序。
torch.manual_seed(12345)
drop_c = {}
for k in dropout_info.keys(): # we don't know the order
p, typ = dropout_info[k]
drop_c[k] = torch.nn.functional.dropout(torch.ones([int(i) for i in typ.shape],
dtype=getattr(torch, typ.dtype), device="cuda"), p=p)*(1-p)
drop_tvm = {n: tensor_to_tvm(t) for n, t in drop_c.items()}
现在我们可以运行前向。
fw_and_cap_compiled_module.set_input('input', inp_tvm[0])
fw_and_cap_compiled_module.set_input('attention_mask', inp_tvm[1])
fw_and_cap_compiled_module.set_input(**model_params_tvm)
fw_and_cap_compiled_module.set_input(**drop_tvm)
fw_and_cap_compiled_module.run()
我们可以将输出与 PyTorch 的输出进行比较：
torch.manual_seed(12345)
pytorch_model.train()
res = pytorch_model(*inp_c)[0]
numpy.abs(fw_and_cap_compiled_module.get_output(0).asnumpy()-res.detach().cpu().numpy()).max()
这给出了 2.1457672e-06。
非常好。让我们也尝试后向。我们生成一个 grad_out，设置所有变量并运行后向模型并运行后向模型：
gr_out_c = torch.randn(res.shape, device="cuda", dtype=res.dtype)
num_captures = len(capture_vars)
num_regular_outputs = len(fw_and_cap_fn_flattened.body.fields) - num_captures
captured_values = {v.name_hint: fw_and_cap_compiled_module.get_output(num_regular_outputs + i) for i, v in enumerate(capture_vars)}
gr_only_compiled_module.set_input(**drop_tvm)
gr_only_compiled_module.set_input(**model_params_tvm)
gr_only_compiled_module.set_input(**captured_values)
gr_only_compiled_module.set_input('gr:out:0', tensor_to_tvm(gr_out_c))
gr_only_compiled_module.run()
在 PyTorch 端，最简单的方法是重新运行前向（记住重置随机种子）并获取梯度。
torch.manual_seed(12345)
pytorch_model.train()
inp_c_rq = [i.requires_grad_() for i in inp_c]
for p in pytorch_model.parameters():
p.requires_grad_()
res = pytorch_model(*inp_c_rq)[0]
grads_pt = torch.autograd.grad(res, inp_c_rq + list(pytorch_model.parameters()), gr_out_c, allow_unused=True)
它奏效了吗？似乎是这样：
for i, g_pt in enumerate(grads_pt):
print(numpy.abs(gr_only_compiled_module.get_output(i).asnumpy() - g_pt.cpu().numpy()).max())
给了我们一个 1e-5 范围内的数字列表。
但是我们想在 PyTorch 中运行一些东西，对吗？
与 PyTorch 的工作方式保持一致，我们首先定义一个 autograd.Function，它手动完成了我们刚才所做的事情：
在 forward 中：
- 生成 dropout 随机值，
- 运行前向，
- 记录后向所需的捕获、输入和 dropout 值。
在 backward 中，运行后向并返回结果（作为 PyTorch 张量）。
这样，我们就得到了一个调用 TVM 的 PyTorch autograd.Function（我们希望有一个小的包装器）。
现在，我们实现获得方法 add_tvm_dispatch(module, sample_inputs) 的目标所需要做的就是追踪模块，从中创建基于 TVM 的 autograd 函数，然后替换前向，如果适用，该前向调用该函数（使用参数），否则回退到通常的前向。Python 的无限动态性使这种黑客行为相对容易。由于所有这些实际上与 TVM 无关，因此我们在这里不赘述（但您可以查看 配套帖子。）
性能
正如我在开头所说，就性能而言，我们还没有完全达到最终目标。在调优任务之后（以及在 HuggingFace BERT + PyTorch JIT 教程中不太真实的推理示例上），我们运行 100 次 TVM 启用的 BertLayer 前向和后向迭代，类似于我们对推理所做的那样。通过 TVM 进行一次迭代需要 6.2 毫秒，而在 PyTorch 上为 1.3 毫秒。
所以我们的模型通过 TVM 运行正常。但它还没有通常的方法快。这里有机会！
更认真地说，我们有两个直接的途径来提高性能：
- 找到一组更好的捕获节点。
- 在 TVM 图上找到优化方法。
就前者的启发式方法而言（请记住，它很可能是 NP 难题，即我相信它是，但我没有制定正式的证明），人们会想要重新进行廉价的计算，最突出的是逐点计算（或者可能是除 matmul 之外的任何东西？）。但这将在另一天进行。
我希望您喜欢本教程，我期待您在 tv@lernapparat.de 上的评论。
致谢
我与 HugingFace 的人以及 Morgan Funtowicz 进行了许多有趣的讨论。此外，TVM 贡献者在审查 TVM 的补丁和论坛上提出了许多好的评论。本教程的创建由 AMD 赞助。
作者
Thomas Viehmann 是 MathInf GmbH 的创始人，MathInf GmbH 是一家位于德国慕尼黑的精品培训和咨询公司，专注于机器学习和 PyTorch。他是 PyTorch 核心开发人员，并与他人合著了 Deep Learning with PyTorch，该书目前可在 PyTorch 网站免费下载。