关于python:TracedModule-更友好的模型表示方案模型训练到部署的桥梁

2次阅读

共计 6000 个字符，预计需要花费 15 分钟才能阅读完成。

作者：曹文刚 | 旷视 MegEngine 架构师

TracedModule 是 MegEngine 中的一种模型格局，用于脱离模型源码对模型进行训练、量化、图手术和模型转换，它是模型训练到部署之间的桥梁。

图 1 从一个一般 Module 生成 TracedModule

TracedModule 产生自一般的 Module，它通过以下两步失去：

运行一次 Module，记录并捕捉模型运行过程中对输出 Tensor 的所有操作，对应图 1 中的 tm.trace_module
通过一个由 5 种指令所形成的“简略”的 high-level IR(intermediate representation) 来形容捕捉到的程序（一般 Module 中的 forward 办法），对应于图 1 中的 SimpleModule.Graph

TracedModule 的实质依然是一个 Module，它与一般 Module 的区别在于：一般 Module 通过用户实现的 forward 办法形容模型运行过程，而 TracedModule 通过 TracedModule IR 来形容模型的运行过程。TracedModule IR 由 python 的根本数据类型以及 Node 和 Expr 形成，其中：Node 用来示意一个 Tensor 或 Module，并记录了 Tensor 或 Module 的一些信息；Expr 用来示意对 Tensor 或 Module 的操作，它的输出和输入都是 Node。

TracedMdoule IR 中的 Expr 共有以下 5 种：

通过以上 5 种 Expr 即可示意绝大部分模型的运行过程。

如前文的介绍，TracedModule IR 是 TracedModule 中的外围数据结构，它用来形容深度学习模型的计算过程，使模型可能脱离源码而存在。不同的深度学习训练框架都有各自的 IR 形容模型，例如：PyTorch 中的 TorchScript，MindSpore 中的 MindIR，以及 onnx 等。这些 IR 大都是一些绝对 low-level 的 IR，在模型源码向这些 IR 转换时经常会产生 python 的层 op 被转换为多个框架底层 op 组合的景象，例如 pytorch 中的 F.Linear 算子在导出到 TorchScript 时可能被导出为 matmul 和 add 的组合。用户在应用 low-levle 的 IR 表白的模型时会有很多的问题，例如：

不理解底层算子用户可能会很难从模型的可视化构造上与模型源码对应
普通用户学习 IR 构造较为艰难，很难对模型进行图手术（批改模型图构造）或优化

这种 low-level 的 IR 表达能力往往更加齐备，相应的也导致 IR 构造极其简单，失去高层语义，使得用户难以做变换和优化，对模型设计者十分不敌对。业界也提出了一些更加 high-level 的 IR 来解决这些问题，比方 torch.fx 和 pnnx 等，这些 IR 都对构造进行了简化，让 IR 的形容模型中的 op 粒度更高，更贴近算法工程师的视角，使得用户学习更简略，解决模型也更容易。

在 MegEngine 中，由多个底层 op 组合成的 python 层 op 更多，例如 “resize”, “relu6”, “softmax” 等，如果间接通过底层 op 表白模型，将会呈现导出的模型构造谁也不意识的窘况。为了解决这些问题，MegEngine 参考 torch.fx 和 TorchScript 计划，改良失去 TracedModule 计划。TracedModule 的 IR 是一个 high-level 的 IR，它所形容的 op 粒度根本与 MegEngine 的 python 层的 op 统一，模型中的 op 粒度与用户视角统一，用户能够很容易地基于 TracedMdoule IR 对模型进行剖析，优化和转换。另外前文提到 TracedMdoule 的实质是一个 Module，用户也能够不便地应用 MegEngine 的模型训练接口对 TracedModule 模型进行训练或参数微调。

TracedModule 全副由 python 层的数据结构形成，trace_module 函数在捕捉用户代码的运行逻辑时仅记录模型中应用的 MegEngine python 层的 function 或 Module，这使得 TracedModule IR 所形容的 op 粒度根本与 MegEngine 的 python 接口统一，即 TracedModule IR 形容的模型是由更加靠近用户视角的高层 op 形成，这使得用户对模型进行一些剖析和优化时更加的容易，例如：

对 MegEngine 的 python 用户更加敌对，相熟 MegEngine python 接口的用户便人造的相熟了由 TracedModule 示意的模型
转换出的模型构造可视化时更为洁净清晰，用户很容易的便可将转换后的模型构造与模型源码对应
对模型进行剖析，优化和向第三方推理框架转换时更容易，比方：模型量化，算子交融，转换器等

更高层 op 的粒度示意使得模型源码转换为 TracedModule 后的模型构造更加洁净清晰，用户很容易的便能够将转换后的模型构造与模型源码进行对应，便于用户对模型进行剖析和调试。

这里以一个罕用的激活函数 relu6 为例，该激活函数在 MegEngine 中的 python 接口如下所示：

def relu6(x):
    relu6 = _get_relu6_op(x.dtype, x.device)
    (x,) = relu6(x)
    return x
 
def _get_relu6_op(....)
    ...
    def relu6(inputs, f, c):
        (inp,) = inputs[0:1]
        max_0 = f("max", inp, c(0))
        min_6 = f("min", max_0, c(6))
        oup = min_6
        (oup_grad,) = yield (oup,)
        ...
    return relu6
...

relu6 函数在 MegEngine 底层实际上是调用了两个算子，别离是模式为 MAX 和 MIN 的 Elemwise 算子，相熟 MegEngine python 源码的同学应该可能从下面的代码中看出 relu6 的前向实现里调用了两个 elemwise 算子，如图 2 所示。

图 2 relu6 的 python 接口和底层实现

如果将一个调用了 relu6 函数导出至由底层 op 所形成的模型，其可视化后果将会如图 3 所示，能够看到 relu6 变成了两个 Elmwise 算子，在这个构造中咱们看不到任何对于 relu6 的信息，不相熟 MegEngine 底层源码的用户面对这样一个模型是比拟懵的。

图 3 可视化由底层算子形成的 relu6

但如果将该模型代码转化至 TracedModule 后，将会失去如图 4 这样一个模型，能够看到 relu6 这个激活函数的信息残缺的存在于 TracedModule 中，并不会被转变为 Elemwise 等其它算子。用户能够容易的从 TracedModule 中找到与模型源码所对应的模块。

图 4 转化到 TracedModule 中的 relu6

相似 relu6 这样的 op 在 MegEngine 中还有很多，例如 leaky_relu、interpolate、conv_transpose2d 等都由多个底层的 op 拼合而成，有些能够从 python 接口的源码中看出其在底层的实现，有些却不太容易看出。能够设想，一个看起来洁净的应用 MegEngine python 层 op 构建的模型代码，在导出为由框架底层 op 形成的模型后，将会呈现模型作者也很难从模型的可视化构造中找到模型某些构造的窘况。但将模型源码导出为由更高层 op 形成的 TracedModule 后，将不会或很少会呈现模型作者不意识可视化出的模型。

将一个 MegEngine 训练出的模型转换至第三方的推理框架进行推理时，经常须要通过图手术对模型构造进行一些批改来满足第三方框架的要求。基于 TracedModule 对模型进行批改是非常容易的，如前文提到 TracedModule 模型中的 op 粒度与用户视角统一，并且形成 TracedModule 的根本数据结构也都是用户相熟的 python 数据结构，只须要理解 TracedModule IR 的根本组件，用户就能够不便的对 TracedModule 所示意的模型运行过程进行批改。

这里以一个罕用于检测模型 Head 模块中 box 分支的操作为例：

F.relu(conv(bbox_subnet) * scale) / stride

其中 conv 是一个一般的卷积，scale 和 stride 是两个常量 Tensor。在 relu(x)∗y 中，当 y >0 时，显然 relu(x)∗y 与 relu(x∗y) 等价，所以在转换上述操作时，经常会将 scale 和 stride 吸到 conv 的权重中，吸掉 scale 和 stride 后的模型构造将更加的简略，并且也不便转换到一些算子较少的两头模型格局，例如 caffe。在 TracedModule 中咱们能够很容易的定位上述操作，并利用图手术接口实现对 scale 和 stride 的排汇。图手术代码如下所示。

graph = traced_head.graph
# 由 conv 的权重排汇 sacle 和 stride
traced_head.conv.weight *= (traced_head.scale / traced_head.stride)
traced_head.conv.bias *= (traced_head.scale / traced_head.stride)
 
# 移除 Graph 中的乘 scale
mul_expr = graph.get_expr_by_id(5).as_unique()
graph.replace_node({mul_expr.outputs[0]: mul_expr.inputs[0]})
# 移除 Graph 中的除 stride
div_expr = graph.get_expr_by_id(8).as_unique()
graph.replace_node({div_expr.outputs[0]: div_expr.inputs[0]})
# 删除 Graph 中无有用的 expr
graph.compile()

如图 5 所示，模型批改完之后，通过 print(graph) 就能间接看到批改之后的图是否满足预期。另外，因为 TracedModule 的 runtime 是 MegEngine 的动态图，在模型运行或图手术时十分的容易调试。

图 5 图手术优化前后的 Head 模块

写到这里可能有人会问，间接批改模型源码之后再转换岂不是更简略？但这会带来另外的问题，比方：模型落地过程中，模型可能会通过好几个人的解决；援用第三方库（例如 basecls 等）进行模型生产时，间接批改底层源码显然是不通用的等。

为了晋升用户体验，TraedMdoule 提供了许多罕用的图手术接口，并尽可能的使用户在应用图手术接口时不须要了解和关注图外部的变动细节。在模型图手术之后，用户能够通过打印 Graph 查看图手术后的图是否合乎预期，也能够像运行一般 Module 一样间接运行 TracedModule 来查看模型输入后果是否正确。可能应用 MegEngine 构建模型的用户，根本在理解 TracedModule 根本组件后，就能够对 TracedMdoule 模型进行图手术。另外，咱们为每一个图手术接口写了具体的应用办法，并提供了一些常见的模型图手术例子供参考，欢送大家来试用。

模型量化是深度学习模型部署过程中的一个重要环节，可能无效缩小模型运行时所占用的计算资源，进步模型的运行速度。各个深度学习训练和推理框架都反对对模型的量化，MegEngine 同样提供了模型量化模块和丰盛的模型量化算法。模型量化的办法大抵能够分为以下两种：

量化感知训练（Quantization Aware Training, QAT），个别是在训练时插入伪量化算子来模仿量化，进而缓减量化带来精度损失
训练后量化（Post-Training Quantization, PTQ），个别是利用无限的输出数据对训练好的模型的权重和激活值进行量化

大多数的推理框架都反对 PTQ 办法对模型进行量化，用户只须要提供浮点模型和输出数据集，个别就能够利用框架提供的量化工具实现模型的量化。然而在 PTQ 无奈满足模型的精度的要求时，便须要借助 MegEngine 等训练框架应用 QAT 办法对模型进行量化，进而进步模型量化的精度。

为了更好的反对 MegEngine 量化训练后的模型部署至第三方推理平台进行推理，MegEngine 团队开发了基于 TracedModule 模型转换工具 mgeconvert 来反对量化模型部署到第三方。TracedModule 不仅反对 MegEngine 底层的量化形式，同时也反对各种自定义的量化算法，这使得基于 TracedModule 导出的量化模型，个别在转换后也可能满足指标平台的量化要求，减小定点模型和伪量化模型之间的差别。用户只须要将 TracedModule 模型输出到 mgeconvert 就能够失去以下两类模型：

浮点模型表示（caffe, onnx, tflite）+ 量化参数文件
定点模型表示（tflite）

即 mgeconvert 既反对导出指标平台的浮点模型和量化参数文件，也反对导出指标平台的定点模型。用户能够不便的应用 MegEengine 量化模块对模型进行量化，量化后也能够不便的应用 mgeconvert 将模型转到预期的推理平台，mgeconvert 应用办法见这里，欢送试用。

TracedModule 是 MegEngine 设计的一种模型格局，设计之初便着重思考了面向用户视角的 op 粒度，模型图手术，量化模型部署等问题，并在文中对这些问题以及 TracedModule 的成果进行了简略的介绍。将来 MegEngine 团队也会开发更多基于 TracedModule 的模型发版工具，例如：模型量化工具，模型优化工具等。最初，欢送大家来试用 TracedModule，也欢送大家提出倡议来一起欠缺 TracedModule。

正文完