TracedModule 基本概念

TracedModule 来源于普通的 Module，但它与普通 Module 不同的是其 TracedModule.forward 方法的执行逻辑通过 InternalGraph 来描述。

下面的例子展示了 Module、TracedModule 以及 InternalGraph 之间的关系。

import megengine.module as M
import megengine.functional as F
import megengine as mge

class SimpleModule(M.Module):
    def __init__(self):
        super().__init__()
        self.linear = M.Linear(4, 5)
        self.param = mge.Parameter([1])

    def forward(self, x):
        x = x + mge.Tensor([1])
        x = F.relu(x)
        return self.linear(x + self.param)

首先创建一个 SimpleModule 实例：

>>> module = SimpleModule()
>>> print(module)
SimpleModule(
  (linear): Linear(in_features=4, out_features=5, bias=True)
)

调用 trace_module 将普通 Module 转换为 TracedModule:

>>> import megengine.traced_module as tm
>>> inp = F.zeros(shape = [3, 4])
>>> traced_module = tm.trace_module(module, inp)
>>> print(traced_module)
TracedModule(
  (linear): Linear(in_features=4, out_features=5, bias=True)
)

traced_module 拥有 graph 属性，graph 描述了 SimpleModule.forward 的执行逻辑：

>>> graph = traced_module.graph
>>> print(graph)
SimpleModule.Graph (self, x) {
        %2:     const_tensor = Constant() -> (Tensor)
        %3:     add_out = x.__add__(const_tensor, )
        %4:     relu_out = nn.relu(add_out, )
        %5:     linear = getattr(self, "linear") -> (Linear)
        %6:     param = getattr(self, "param") -> (Tensor)
        %7:     add_out_1 = relu_out.__add__(param, )
        %8:     linear_out = linear(add_out_1, )
        return linear_out
}

我们可以看到 SimpleModule.Graph(self, x) 与 SimpleModule.forward(self, x) 的代码逻辑一致。

一个普通的 Module 可通过 trace_module 方法将其转换为 TracedModule。 在转换过程中，用户自定义的 Module 将被转换为 TracedModule，内置 Module（如 Linear, Conv2d 等）不作转换。

graph 是 TracedModule 中最重要的属性，其实际是一个 InternalGraph, 构成 InternalGraph 的基本单元为 Node 和 Expr。

Node

Node 的常用属性以及方法的使用例子请参考 TracedModule 的常用方法。

通过 Node 来描述 forward 中的 Tensor 或 Module。

Class Node:
    expr : Expr # 描述了该 Node 由哪个 Expr 生成
    users : List[Expr] # 描述了该 Node 被哪些 Expr 使用

    @property
    def top_graph(self) -> InternalGraph: ... # 该 Node 所属的 InternalGraph
    @property
    def name(self) -> str: ... # 该 Node 的名字
    @property
    def qualname(self) -> str: ... # 生成该 Node 的 Module 的名字

Node 的 expr 属性记录了生成该 Node 的 Expr, 例如 SimpleModule 的输出是调用 linear 这个子 module 生成的。

>>> out_node = graph.outputs[0]
>>> print(out_node)
linear_out
>>> out_node.expr
%8:     linear_out = linear(add_out_1, )

Node 的 users 属性记录了该 Node 的被哪些 Expr 作为输入, 例如 SimpleModule 的中的输入 x 是 x = x + 1 这个的输入。

>>> inp_node = graph.inputs[1]
>>> print(inp_node)
x
>>> inp_node.users
[%3:    add_out = x.__add__(const_tensor, )]

Node 的 name 属性是该 Node 的名字，该名字在其所属的 graph 中是唯一的。

Node 的 qualname 属性记录了该 Node 是由哪个 Module 中所生成的，可以从 qualname 得到该 Module 的名字。 例如，out_node.qualname 为 ‘SimpleModule.linear.[out]’, 表示 out_node 是 SimpleModule 中 linear 这个子 module 的输出。

>>> out_node = graph.outputs[0]
>>> out_node.qualname
'SimpleModule.linear.[out]'

InternalGraph 中的 Node 有两种：

• TensorNode：描述一个 Tensor，记录了该 Tensor 的 dtype 、shape 和 qparams 等信息

  >>> x = graph.inputs[1]
  >>> type(x)
  <class 'megengine.traced_module.node.TensorNode'>
  >>> x.shape
  (3, 4)
  >>> x.dtype
  numpy.float32
  

• ModuleNode：描述一个 Module，记录了该 Module 的类型，以及对应的 Module

  >>> self = graph.inputs[0]
  >>> type(self)
  <class 'megengine.traced_module.node.ModuleNode'>
  >>> x.owner # 通过 owner 属性访问该 ModuleNode 所对应的 Module
  TracedModule(
  (linear): Linear(in_features=4, out_features=5, bias=True)
  )
  

Expr

Expr 的常用属性以及方法的使用例子请参考 TracedModule 的常用方法。

通过 Expr 来描述一个 forward 中的某个表达式。 一个 Expr 由表达式的输入 ( inputs )、 输出 ( outputs )、 以及由输入到输出的执行逻辑 ( interpret ) 构成。

Class Expr:
    inputs : List[Node] # 输入的 Node
    const_val : List[int,float,...] # 输入的常量
    outputs : List[Node] # 输出的 Node

    @property
    def top_graph(self) -> InternalGraph:... # 该 Expr 所属的 InternalGraph

    def interpret(self, *args, **kwargs):... # 根据输入执行该 expr

Expr 的子类分别有：

• GetAttr: 获取 TracedModule 的中的某个属性，该 Expr 保存一个 name 字符串（用来描述要获取的属性）， 接受一个 ModuleNode 作为输入，它的执行逻辑为 outputs = getattr(inputs[0], name)。

  例如：SimpleModule.forward 中的 self.param 将会被解释为 %6: param = getattr(self, "param") -> (Tensor)， self.linear 将会被解释为 %5: linear = getattr(self, "linear") -> (Linear)，这两个 GetAttr 的输入均为 self 这个 ModuleNode。

  >>> exprs = graph.exprs(recursive=False).aslist()
  >>> exprs[6]
  %6:    param = getattr(self, "param") -> (Tensor)
  >>> exprs[6].inputs
  [self]
  >>> exprs[6].outputs
  [param]
  

• CallMethod: 调用变量（Module 或 Tensor）的一个方法，该 Expr 保存一个 method 字符串（用来描述调用变量的哪个方法）， 接受多个输入（第一个输入为变量本身，即 self）。它的执行逻辑为 otuputs = getattr(inputs[0], method)(\*inputs[1:])。

  例如：SimpleModule.forward 中的 x = x + self.param 将会被解释为 %7: add_out_1 = relu_out.__add__(param, )， 这个 expr 是指调用了 x 的 __add__ 方法，输入为 x 和 param。

  >>> exprs = graph.exprs(recursive=False).as_dict()
  >>> exprs[7]
  %7:     add_out_1 = relu_out.__add__(param, )
  >>> exprs[7].inputs
  [relu_out, param]
  >>> exprs[7].outputs
  [add_out_1]
  

• CallFunction: 调用 megengine 内置的某个函数，该 Expr 保存一个 func，接受多个输入。 它的执行逻辑为 outputs = func(\*inputs) 。

  例如：SimpleModule.forward 中的 x = F.relu(x) ，将会被解释为 %4: relu_out = nn.relu(add_out, ), 表示调用了 nn.relu 这个 function，其输入为 add_out。

  >>> exprs = graph.exprs(recursive=False).as_dict()
  >>> exprs[4]
  %4:    relu_out = nn.relu(add_out, )
  >>> exprs[4].inputs
  [add_out]
  >>> exprs[4].outputs
  [relu_out]
  

• Constant: 产生一个常量，该 Expr 会记录一个不会改变的 value（Module 或 Tensor），不接受输入，它的执行逻辑为 outputs = value。

  例如：SimpleModule.forward 中的 mge.Tensor([1]) 将会被解释为 %2: const_tensor = Constant() -> (Tensor)， 表示一个生成常量 Tensor。

  >>> exprs = graph.exprs(recursive=False).as_dict()
  >>> exprs[2]
  %4:    relu_out = nn.relu(add_out, )
  >>> exprs[2].inputs
  []
  >>> exprs[2].outputs
  [const_tensor]
  

• Input: 表示 Module.forward 的输入，仅仅是一个占位符的作用。真正推理的时候会将其替换为真正的 Tensor。

所有的 Node 在实际执行推理的时候（interpret）都会被替换为实际的 Tensor 或者 Module。

InternalGraph

InternalGraph 的常用属性以及方法的使用例子请参考 TracedModule 接口介绍。

将 Module.foward 中的每一条语句都解释为由 Expr 组成的执行序列就构成了最终的 InternalGraph。

Class InternalGraph:
    _exprs : List[Expr]

    def interpret(self, *inputs):...

    @property
    def inputs(self):...

    @property
    def outputs(self):...

InternalGraph 包含以下三个属性：

• _exprs: 按执行顺序排列的 Expr 列表

  >>> graph._exprs
  [%2:    const_tensor = Constant(<class 'megengine.tensor.Tensor'>) -> (Tensor),
   %3:    add_out = x.__add__(const_tensor, ),
   %4:    relu_out = nn.relu(add_out, ),
   %5:    linear = getattr(self, "linear") -> (Linear),
   %6:    param = getattr(self, "param") -> (Tensor),
   %7:    add_out_1 = relu_out.__add__(param, ),
   %8:    linear_out = linear(add_out_1, )]
  

• inputs: 该 graph 的输入 Node

  >>> graph.inputs
  [self, x]
  

• outputs: 该 graph 的输出 Node

  >>> graph.outputs
  [linear_out]
  

在解析 Module.forward 的过程中，会将 forward 里的每一个执行语句描述为 Expr，并按执行次序依次添加到 _exprs 属性里。

在真正推理时，只需要遍历 _exprs 并依次 interpret 即可得到与执行原 Module 一样的结果。