TracedModule 常见图手术#

Note

阅读本文所展示的图手术例子，需要先了解 TracedMdoule 图手术的相关接口和用法，请参考 Node 和 Expr 的查找方法和图手术常用方法。

修改 Node 的名字#

修改 graph 中的 Node 的名字，可以直接对 Node.name 赋值即可，但在赋值时要确保新的名字未被 graph 中其它的 Node 所使用。

例如修改某 graph 第一个输出的 Node 的名字，可以通过直接获取 graph 的 outputs，来获得输出 Node，再直接重新设置 Node 的 name 就可对其重命名。

>>> out_node = traced_net.graph.outputs[0]
>>> out_node.name = "I_am_output"

为模型添加前后处理#

由于 TracedModule 可以被重新 trace，因此在加前后处理时，可以新写一个 Module ，并将前处理，主体模型和后处理作为新 Module 的 3 个子 Module，并在新 module 的 forward 函数中分别调用 3 个 module 就完成了前后处理的添加。例子如下：

将一些常量吸收到卷积里#

对于一些基于 anchor 的检测算法，经常会在卷积的输出后，对卷积结果乘 stride 或除 anchor_size，在推理部署时，可以将这些常量吸收到卷积里，基于 TracedModule 可以较容易的实现这些转换，如下面的例子：

将一些 OP 转换为 fp16#

对于一些计算量特别大的全连接层，会占用较多的存储资源，可以通过将其转换为 fp16 计算减少其占用的资源, 如下面的例子：

通过 Graph 确定数据流向#

在量化训练时，常常会对 concat 的输入做某些约束，通过 TracedModule 可以轻易的找到这些 concat 的输入是来自于哪个内置的 function 或 Module 的输出，如下面的例子。

Conv 和 BN 融合#

在推理或量化训练时，常常需要将 Conv 和 Bn 融合到一起，基于 TracedModule 的 Graph 可以找到满足融合条件的 Conv 和 Bn，并以图手术的方式将其融合，如下面的例子。

将 BN 融合到 Conv 中

import numpy as np
import pickle

import megengine.functional as F
import megengine.module as M
import megengine.module.qat as Q
import megengine.traced_module as tm
from megengine.traced_module.expr import CallMethod
from megengine.traced_module.node import ModuleNode
import megengine as mge

class Net(M.Module):
    def __init__(self,):
        super().__init__()
        self.conv = M.Conv2d(3,16,1, bias=False)
        self.bn = M.BatchNorm2d(16)
        self.conv_bn = M.Sequential(
            M.Conv2d(16,16,1,bias=False),
            M.BatchNorm2d(16)
        )

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = F.relu(x)
        x = self.conv_bn(x)
        return x

def _fuse_conv_bn(conv : M.Conv2d, bn : M.BatchNorm2d = None):
    weight, bias = conv.weight, conv.bias
    target_cls = M.ConvBn2d
    if not conv.training:
        class FakeCls:
            def __init__(self, conv, bn):
                self.conv = conv
                self.bn = bn
            def apply_quant_weight(self, inp):
                return inp

        weight, bias = Q.ConvBn2d.fold_weight_bias(
            FakeCls(conv, bn),
            bn.running_mean,
            bn.running_var
        )
        target_cls = M.Conv2d
    this_module = target_cls(
        conv.in_channels,
        conv.out_channels,
        conv.kernel_size,
        conv.stride,
        conv.padding,
        conv.dilation,
        conv.groups,
        conv.bias is not None,
        conv.conv_mode,
        conv.compute_mode,
        name=conv.name,
    )
    if conv.training:
        this_module.conv.weight = weight
        this_module.conv.bias = bias
        this_module.bn = bn
    else:
        this_module.weight = weight
        this_module.bias = bias
    return this_module

def fuse_bn_transform():
    net = Net()
    inp = mge.Tensor(np.random.random(size = (1,3, 16, 16)), dtype=np.float32)

    traced_net = tm.trace_module(net, inp)
    obj = pickle.dumps(traced_net)
    traced_net = pickle.loads(obj)

    graph = traced_net.graph

    for conv_node in graph.get_module_by_type(M.Conv2d).as_list():
        if len(conv_node.users) > 1:
            continue
        conv_expr = conv_node.users[0]
        conv_out_node = conv_expr.outputs[0]

        if len(conv_out_node.users) > 1:
            # conv -> bn，conv 的输出只能被 bn 使用
            continue

        # 判断 conv 之后的 expr 是否是 bn
        bn_expr = conv_out_node.users[0]
        if not isinstance(bn_expr, CallMethod):
            continue
        bn_node = bn_expr.inputs[0]
        if not isinstance(bn_node, ModuleNode) or bn_node.module_type != M.BatchNorm2d:
            continue

        conv_module = conv_node.owner
        bn_module = bn_node.owner

        new_module = _fuse_conv_bn(conv_module, bn_module)

        cur_graph = conv_node.top_graph
        self_node = cur_graph.inputs[0]
        self_module = self_node.owner
        name = conv_module._name

        # 将 fuse 后的 module 设置到 调用 conv 的 module 上
        setattr(self_module, conv_module._name, new_module)
        inp_node = conv_expr.inputs[1]
        bn_out_node = bn_expr.outputs[0]

        # 将 fuse 后的 module 以图手术的方式 insert 到 graph 中
        with cur_graph.insert_exprs():
            fused_conv_out = getattr(self_node, name)(inp_node)

        cur_graph.replace_node({bn_out_node: fused_conv_out})
        cur_graph.compile()

    gt = net(inp)
    actual = traced_net(inp)
    np.testing.assert_allclose(gt.numpy(), actual.numpy(), atol=5e-2)


if __name__ == "__main__":
    fuse_bn_transform()