作者：曹文刚|MegEngineArchitectTracedModule简介TracedModule是MegEngine中的一种模型格式，用于在没有模型源代码、图运算和模型转换的情况下对模型进行训练和量化，是模型训练和部署之间的桥梁。图1GenerateTracedModulefromanordinaryModuleTracedModule是从一个普通的Module生成的，通过以下两步得到：运行一次Module，记录并捕获模型运行过程中对输入Tensor的所有操作，对应tm.trace_module图1中通过一个由5种指令组成的“简单”高级IR（中间表示）来描述捕获的程序（commonModule中的forward方法），对应图1中的SimpleModule.GraphTracedModule。本质仍然是一个Module，它与普通Module的区别在于，普通Module通过用户实现的forward方法来描述模型的运行过程，而TracedModule通过TracedModuleIR来描述模型的运行过程。TracedModuleIR由python的基本数据类型和Node、Expr组成，其中：Node用于表示一个Tensor或Module，记录了Tensor或Module的一些信息；Expr用来表示Tensor或者Module的运算，它的输入和输出都是Node。TracedMdouleIR中有五种Expr：|原稿|含义|范例||---|----|--------------||输入|放置角色|\||常量|意味着生成一个常量Tensor|mge.Tensor([1])->%2:const_tensor=Constant()->(张量)||获取属性|表示获取Module的属性|self.linear->%5:linear=getattr(self,"linear")->(线性)||调用方法|表示调用Module的forward方法或者Tensor的某个方法|x+self.param->%7:add\_out\_1=relu\_out.\__add\_\_(param,)||调用函数|调用函数|F.relu(x)->%4:relu\_out=nn.relu(add\_out,)|大多数模型的运行过程都可以用以上五种Expr表示。为什么是TracedModule？上文提到，TracedModuleIR是TracedModule中的核心数据结构，用于描述深度学习模型的计算过程，使模型可以无源代码存在。不同的深度学习训练框架都有自己的IR描述模型，例如：PyTorch中的TorchScript、MindSpore中的MindIR、onnx等，这些IR大多是比较底层的IR。模型源码在转换成这些IR时，经常会出现python层op转换成多个框架底层op组合的情况。例如，pytorch中的F.Linear运算符被导出可能作为matmul和add的组合被导出到TorchScript。用户在使用低级IR表达的模型时会遇到很多问题。例如，不了解底层算子的用户可能很难从模型源码对应的模型可视化结构中学习IR结构。对模型进行图外科手术（修改模型图结构）或优化低层IR表达能力往往更完整，这也导致极其复杂的IR结构和高层语义的丢失，使用户难以理解进行改造和优化。非常不友好。业界也提出了一些更高层次的IR来解决这些问题，例如torch.fx和pnnx等，这些IR简化了结构，使得IR描述模型中的op粒度更高，更贴近算法工程师。透视图使用户更容易学习和处理模型。在MegEngine中，更多的是由多个底层op组成的python层op，比如“resize”、“relu6”、“softmax”等，如果直接通过底层op来表达模型，就没有了一个在导出模型结构中知道的困境。为了解决这些问题，MegEngine参考了torch.fx和TorchScript方案，改进了TracedModule方案。TracedModule的IR是高级IR。它描述的op粒度与MegEngine的python层的op粒度基本一致。模型中的操作粒度与用户的视角一致。用户可以基于TracedMdouleIR轻松分析模型。优化和转换。另外，如前所述，TracedMdoule的本质是一个Module。用户也可以方便地使用MegEngine的模型训练接口对TracedModule模型的参数进行训练或微调。TracedModule有什么好处？TracedModule都是由python层的数据结构组成的。trace_module函数在捕获用户代码的运行逻辑时，只记录了模型中使用的MegEnginepython层的函数或Module，这使得TracedModuleIR描述的op粒度与MegEngine的python接口基本一致，即，TracedModuleIR描述的模型由更贴近用户视角的high-levelops组成，方便用户分析和优化模型。自然熟悉TracedModule所代表的模型，转换后的模型结构在可视化时更加干净清晰。用户可以方便地分析转换后的模型结构和模型对应的模型源码，优化并转换为第三方推理框架，例如模型量化、算子融合、转换器等干净的模型表示。更高级ops的粒度表示使得模型源代码转换为TracedModule后模型结构更清晰。用户可以方便的将转换后的模型结构与模型源码对应起来，方便用户分析调试模型。这里以一个常用的激活函数relu6为例。MegEngine中这个激活函数的python接口如下：defrelu6(x):relu6=_get_relu6_op(x.dtype,x.device)(x,)=relu6(x)returnxdef_get_relu6_op(....)...defrelu6(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,)...returnrelu6...relu6函数实际上调用了MegEngine底层的两个算子，分别是模式为MAX和MIN的Elemwise算子。熟悉MegEnginepython源码的同学应该能从上面的代码看出relu6的前向实现中调用了两个elemwise算子，如图2所示。图2relu6的Python接口和底层实现如果一个函数callsrelu6被导出为底层ops组成的模型，可视化结果如图3所示。可以看到relu6变成了两个Elmwise算子，在这个结构中，我们看不到relu6的任何信息。不熟悉MegEngine底层源码的用户，在面对这样的模型时会比较迷茫。图3可视化了由底层算子组成的relu6，但是如果将模型代码转换为TracedModule，会得到如图4这样的模型。可以看到relu6的激活函数信息完全存在于TracedModule中，不会转换为Elemwise等其他算子。用户可以很容易的从TracedModule中找到模型源码对应的模块。图4Relu6转化为TracedModuleMegEngine中有很多类似relu6的ops，比如leaky_relu、interpolate、conv_transpose2d等都是由多个底层ops组成的，其中一些可以从python接口的源码中看到。一些底层实现并不容易看。可以想象，使用MegEnginepython层的op，一个看起来很干净的模型代码，导出为由框架底层op组成的模型，模型作者很难从中找到模型的结构模型的视觉结构。尴尬。但是，将模型源码导出为由更高层ops组成的TracedModule后，模型作者不认识的可视化模型将没有或很少。直观的模型图外科在将MegEngine训练的模型转换为第三方推理框架进行推理时，往往需要通过图外科对模型结构进行修改，以满足第三方框架的要求。基于TracedModule修改模型非常容易。如前所述，TracedModule模型中的op粒度与用户的角度是一致的，构成TracedModule的基本数据结构也是用户熟悉的python数据结构。您只需要了解TracedModuleIRComponent的基础知识，用户就可以轻松修改TracedModule代表的模型的运行流程。这里举一个模型head模块中经常用来检测box分支的操作??的例子：F.relu(conv(bbox_subnet)*scale)/stride其中conv是一个普通的卷积，scale和stride是两个常数张量。在relu(x)?y中，当y>0时，显然relu(x)?y等价于relu(x?y)，所以在转换上述操作时，scale和stride往往被吸收到conv的权重中，吸收scale和stride后的模型结构会更简单，也方便转换成一些算子较少的中间模型格式，比如caffe。在TracedModule中，我们可以很方便的定位到上述操作，通过图形化的操作界面，完成对scale和stride的吸收。图操作代码如下所示。graph=traced_head.graph#通过convtraced_head.conv.weight*=(traced_head.scale/traced_head.stride)traced_head.conv.bias*=(traced_head.scale/traced_head.stride)的权重吸收sacle和stride#从Graph中删除scalemul_expr=graph.get_expr_by_id(5).as_unique()graph.replace_node({mul_expr.outputs[0]:mul_expr.inputs[0]})#去掉Graph中的除法stridediv_expr=graph.get_expr_by_id(8)。as_unique()graph.replace_node({div_expr.outputs[0]:div_expr.inputs[0]})#删除在Graph中没有用的exprgraph.compile()。如图5所示，模型修改后，print(graph)可以直接看到修改后的图是否符合预期。另外，由于TracedModule的runtime是MegEngine的一个动态图，在模型运行或者图术的时候非常容易调试。Figure5图5Headmodule优化前后可能有人会问，直接修改模型源码再转换不是更方便吗？但这会带来其他问题，例如：在模型落地的过程中，模型可能会被多人处理；在引用第三方库（如basecls等）进行模型制作时，直接修改底层源码等显然不通用。为了提升用户体验，TraedMdoule提供了很多常用的图外科接口，并尽量让用户在使用图外科界面时无需了解和关注图内变化的细节。模型图运行后，用户可以打印Graph来检查图运行后的图是否符合预期，也可以像运行普通Module一样直接运行TracedModule来检查模型输出是否正确。可以使用MegEngine搭建模型的用户，在基本了解TracedModule的基本组件后，可以对TracedMdoule模型进行图手术。另外，我们为每个图外科接口都写了详细的使用方法，并提供了一些常用的模型图外科实例供参考，欢迎大家试用。便捷的量化模型部署模型量化是深度学习模型部署过程中的重要环节，可以有效降低模型运行时占用的计算资源，提高模型的运行速度。所有的深度学习训练和推理框架都支持模型量化，MegEngine也提供了模型量化模块和丰富的模型量化算法。模型量化的方法大致可以分为以下两种：量化感知训练（QuantizationAwareTraining，QAT），一般在训练时插入伪量化算子来模拟量化，从而减少量化带来的精度损失。训练后量化（Post-TrainingQuantization，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的建议一起。
