InternalGraph

class InternalGraph(name, qualname)

InternalGraph 是 TracedModule 中使用的主要数据结构。它用来表示 Module 的 forward 方法的执行过程。

例如，下面的代码

import megengine.random as rand
import megengine.functional as F
import megengine.module as M

import megengine.traced_module as tm

class MyModule(M.Module):
    def __init__(self):
        super().__init__()
        self.param = rand.normal(size=(3, 4))
        self.linear = M.Linear(4, 5)

    def forward(self, x):
        return F.relu(self.linear(x + self.param))

net = MyModule()

inp = F.zeros(shape = (3, 4))
traced_module = tm.trace_module(net, inp)

Will produce the following InternalGraph:

print(traced_module.graph)

MyModule.Graph (self, x) {
        %2:     linear = getattr(self, "linear") -> (Linear)
        %3:     param = getattr(self, "param") -> (Tensor)
        %4:     add_out = x.__add__(param, )
        %5:     linear_out = linear(add_out, )
        %6:     relu_out = nn.relu(linear_out, )
        return relu_out
}

add_input_node(shape, dtype='float32', name='args')

向计算图中添加一个输入节点。

新 Node 将是最后一个位置参数。

参数:

• shape (Tuple[int]) – 新输入 Node 的形状。

• dtype (str) – 新输入 Node 的 dtype 。默认值：float32

• name (str) – 新输入 Node 的名字。当图中使用该名字时，如果重名将为其添加后缀。

add_output_node(node)

向计算图中添加一个输出节点。

调用 add_output_node 后，图形输出将变成一个 tuple 。tuple 的第一个元素是原始输出，第二个元素是 node 。

例如，下面的代码

import megengine.functional as F
import megengine.module as M
import megengine.traced_module as tm

class MyModule(M.Module):
    def forward(self, x):
        x = x + 1
        return x

net = MyModule()

inp = F.zeros(shape = (1, ))
traced_module = tm.trace_module(net, inp)
graph = traced_module.graph
inp_node = graph.inputs[1]
out_node = graph.outputs[0]
graph.add_output_node(inp_node)
graph.add_output_node(out_node)
out = traced_module(inp)

将生成以下 InternalGraph 和 out:

print(graph)
print(out)

MyModule.Graph (self, x) {
        %2:     add_out = x.__add__(1, )
        return add_out, x, add_out
}
((Tensor([1.], device=xpux:0), Tensor([0.], device=xpux:0)), Tensor([1.], device=xpux:0))

compile()

删除未使用的 expr 。

eval(*inputs)

调用这个方法来执行计算图。

参数:

inputs (Tuple[Tensor]) – 对应于 graph.inputs[1:] 的向量。

exprs(recursive=True)

获取构成这个计算图的 Exprs。

参数:

recursive – 是否获取子图中的 Exprs 。默认值：True

返回:

ExprFilter 包含计算图的所有 Exprs.

get_dep_exprs(nodes)

获取依赖 nodes 的 Exprs.

参数:

nodes (Sequence[Node]) – 一个 Node 列表。

返回类型:

List[Expr]

返回:

一个相关的 Expr 列表。

get_expr_by_id(expr_id=None, recursive=True)

通过 id 来过滤 Exprs.

参数:

• expr_id (Optional[List[int]]) – 一个 int 列表。

• recursive – 是否获取子图中的 Exprs 。默认值：True

返回:

一个 ExprFilterExprId.

get_function_by_type(func=None, recursive=True)

通过 CallFunction 的类型来过滤 Exprs.

参数:

• func (Optional[Callable]) – 一个内置函数，例如 F.relu.

• recursive – 是否获取子图中的 Exprs 。默认值：True

返回:

一个 ExprFilterCallFunction.

get_method_by_type(method=None, recursive=True)

通过 CallMethod 的类型来过滤 Exprs.

参数:

• method (Optional[str]) – 一个方法的字符串，例如 “__add__”.

• recursive – 是否获取子图中的 Exprs 。默认值：True

返回:

一个 ExprFilterCallMethod.

get_module_by_type(module_cls, recursive=True)

通过 ModuleNode 的 module_type 来过滤 Nodes.

参数:

• module_cls (Module) – Module 的子类。

• recursive – 是否获取子图中的 Nodes. 默认值：True

返回:

一个 NodeFilterType.

get_node_by_id(node_id=None, recursive=True)

通过 id 来过滤 Nodes.

Node 的 id 可以通过下面的代码获得

# node : Node
print("{:i}".format(node))
print(node.__format__("i"))
# graph : InternalGraph
print("{:i}".format(graph))
print(graph.__format__("i"))

参数:

• node_id (Optional[List[int]]) – 一个 int 列表。

• recursive – 是否获取子图中的 Nodes. 默认值：True

返回:

一个 NodeFilterNodeId.

get_node_by_name(name=None, ignorecase=True, recursive=True)

通过全名来过滤 Nodes.

Node 的全名可以通过下面的代码获得

# node : Node
print("{:p}".format(node))
print(node.__format__("p"))
# graph : InternalGraph
print("{:p}".format(graph))
print(graph.__format__("p"))

参数:

• name (Optional[str]) – 一个在 glob 语法中可以包含 ? 和 * 来匹配单个或任意字符的字符串。

• ignorecase (bool) – 是否忽略这种情况。默认值：True

• recursive – 是否获取子图中的 Nodes. 默认值：True

返回:

一个 NodeFilterName.

property inputs

获取此计算图的输入节点列表。

返回类型:

List[Node]

返回:

一个 Node 列表。

insert_exprs(expr=None)

初始化 trace 模式和插入位置。

当在 with 语句中使用时，将临时设置 trace 模式，然后在with语句退出时恢复正常模式：

with graph.insert_exprs(e): # set the trace mode
    ... # trace function or module
... # inert exprs into graph and resotre normal mode

参数:

expr (Optional[Expr]) – 在哪个 expr 之后插入。如果为None，插入位置将根据输入节点自动设置。

返回:

一个资源管理器，它将在 __enter__ 里初始化 trace 模式，并在 __exit__ 里恢复正常模式。

property name

获取这个计算图的名称。

返回类型:

str

nodes(recursive=True)

获取构成这个计算图的节点。

参数:

recursive – 是否获取子图中的 Nodes. 默认值：True

返回:

包含此计算图的所有节点的 NodeFilter 。

property outputs

获取此计算图的输出节点列表。

返回类型:

List[Node]

返回:

一个 Node 列表。

property qualname

获取此计算图的 qualname 。qualname 可以用来从被跟踪的模块或普通模块中获取子模块。

示例

import megengine.module as M
import megengine.traced_module as tm
import megengine as mge

class block(M.Module):
    def __init__(self):
        super().__init__()
        self.relu = M.ReLU()

    def forward(self, x):
        return self.relu(x)

class module(M.Module):
    def __init__(self):
        super().__init__()
        self.block = block()

    def forward(self, x):
        x = self.block(x)
        return x

net = module()
traced_net = tm.trace_module(net, mge.Tensor([0.]))

qualname = traced_net.block.graph.qualname  # qualname = "module.block"
qualname = qualname.split(".", 1)[-1]  # qualname = "block"

assert qualname in list(map(lambda x: x[0], net.named_modules()))
assert qualname in list(map(lambda x: x[0], traced_net.named_modules()))

返回类型:

str

replace_node(repl_dict)

替换图中的节点。

参数:

repl_dict (Dict[Node, Node]) – 指定如何替换节点的 map {old_Node: new_Node} 。

reset_outputs(outputs)

重置计算图的输出节点。

备注

此方法只支持重置没有父图的计算图的输出。

参数:

outputs – 内部元素为 Node 的对象。支持元组、列表、字典等。

例如，下面的代码

import megengine.functional as F
import megengine.module as M
import megengine.traced_module as tm

class MyModule(M.Module):
    def forward(self, x):
        x = x + 1
        return x

net = MyModule()

inp = F.zeros(shape = (1, ))
traced_module = tm.trace_module(net, inp)
graph = traced_module.graph
inp_node = graph.inputs[1]
out_node = graph.outputs[0]
graph.reset_outputs((out_node, {"input": inp_node}))
out = traced_module(inp)

将生成以下 InternalGraph 和 out:

print(graph)
print(out)

MyModule.Graph (self, x) {
        %2:     add_out = x.__add__(1, )
        return add_out, x
}
(Tensor([1.], device=xpux:0), {'input': Tensor([0.], device=xpux:0)})

property top_graph

得到这个计算图的父图。

返回:

一个 InternalGraph.