快速上手 TracedModule¶
TracedModule 本质是一个 Module, 其由普通的 Module 转换而来。 本文以 shufflenet_v2 为例，讲解如何使用 TracedModule 进行模型开发与部署。
Module 转换为 TracedModule¶
本文使用的 shufflenet_v2 模型为 MegEngine 预训练的模型，来自 模型中心 。
我们利用 trace_module 方法将 shufflenet 这个 Module 转换 TracedModule.
import numpy as np
import megengine.functional as F
import megengine.module as M
import megengine as mge
import megengine.traced_module as tm
shufflenet = mge.hub.load("megengine/models", "shufflenet_v2_x1_0", pretrained=True)
shufflenet.eval()
data = mge.Tensor(np.random.random([1, 3, 224, 224]).astype(np.float32))
traced_shufflenet = tm.trace_module(shufflenet, data)
trace_module 方法所返回的 traced_shufflenet 是一个 TracedModule, 但其本质是一个 Module.
>>> isinstance(traced_shufflenet, M.Module)
True
我们可以像普通 Module 一样使用 TracedModule, 例如将其作为一个 Module 的子 Module.
class MyModule(M.Module):
def __init__(self, backbone: M.Module):
super().__init__()
self.backbone = backbone
def forward(self, x):
x = self.backbone(x)
return x
new_net = MyModule(traced_shufflenet)
对于任意一个 Module 我们都可以将其转换为 TracedModule, 包括 TracedModule 也可以被再次转换
（再次 trace 后的模型结构不会发生变化）。
>>> traced_shufflenet = tm.trace_module(traced_shufflenet, data)
>>> traced_new_net = tm.trace_module(new_net, data)
TracedModule 序列化¶
注解
推荐使用
.tm作为 TracedModule 序列化文件的后缀名推荐使用
megengine.save和megengine.load保存和加载 TracedModule
可以直接使用 MegEngine 提供的序列化接口 megengine.save 将 TracedModule 模型序列化到文件中。
>>> mge.save(traced_shufflenet, "traced_shufflenet.tm")
也可以直接使用 python 内置的 pickle 模块将 TracedModule 序列化到文件中。
>>> with open("traced_shufflenet.tm", "wb") as f:
... pickle.dump(traced_shufflenet, f)
对应的，可以分别使用 megengine.load 或 pickle.load 将序列化的 TracedModule 再加载回来。
在脱离模型源码的环境中加载得到的 TracedModule, 依然可以被正常解析与运行。
>>> traced_resnet = mge.load("traced_shufflenet.tm")
>>> with open("traced_shufflenet.tm", "rb") as f:
... traced_resnet = pickle.load(f)
TracedModule 图手术¶
TracedModule 提供了一些方便的图手术接口来修改 TracedModule 的执行逻辑。 图手术的接口可以直接阅读 图手术常用方法, 每一个接口下都提供了如何使用该接口的例子。
我们提供了一些常见的图手术例子在 TracedModule 常见图手术 中，可以了解如何完成对 TracedMdoule 执行逻辑的的修改。
同时，我们提供了一些内置的图手术实现来优化模型结构，包括：
FuseConvBn：将 BatchNorm 融合到 Convolution 中
FuseAddMul：融合连续的常量加法或常量乘法
BackwardFoldScale：将卷积之后的常量乘法融合到卷积中
使用这些优化的接口为 tm.optimize, 具体用法请参考 TracedMdoule 内置模型优化.
TracedModule 模型部署¶
使用 MegEngine 进行部署¶
TracedModule 本质是一个 Module, 使用 MegEngine 进行模型部署与普通 Module 部署方法一致， 可参考 trace & dump 将模型转为 c++ 静态图模型。
以 shufflenet_v2 为例，我们可以直接使用上面被转换为 TracedModule 的 shufflenet_v2 模型 traced_shufflenet,
也可以在脱离 shufflenet_v2 源码的环境中直接加载被序列化后并保存的 traced_shufflenet.tm.
>>> traced_shufflenet = mge.load("traced_shufflenet.tm")
然后调用 trace 方法将 traced_shufflenet 转换为 MegEngine 静态图，
调用 dump 方法将静态图保存为 c++ 模型，同时在 dump 时可以开启 Inference 相关的优化。
import megengine.jit as jit
@jit.trace(symbolic=True, capture_as_const=True)
def fun(data, *, net):
pred = net(data)
return pred
data = mge.Tensor(np.random.random([1, 3, 224, 224]).astype(np.float32))
fun(data, net=traced_shufflenet)
fun.dump("shufflenet_v2.mge", arg_names=["data"], optimize_for_inference=True)
将模型序列化为 c++ 模型并保存后，可以参考 MegEngine Lite C++ 部署模型快速上手 在 c++ 环境中进行模型部署， 或参考 MegEngine Lite Python 部署模型快速上手 在 python 环境中进行模型部署。
使用第三方平台部署¶
对于第三方平台部署，我们开发了基于 TracedModule 的转换器 mgeconvert ， 可以方便地将 TracedModule 所描述的模型结构和参数转换至其它推理框架，例如：caffe, tflite 和 onnx 等， 未来也将支持更多第三方推理框架的转换。